diff --git a/client/src/MainLayout/index.jsx b/client/src/MainLayout/index.jsx index 7f47cc1..dac5831 100644 --- a/client/src/MainLayout/index.jsx +++ b/client/src/MainLayout/index.jsx @@ -202,7 +202,7 @@ export const MainLayout = ({ const debugModeOn = Boolean(window.localStorage.$ANNOTATE_DEBUG_MODE && state) const nextImageHasRegions = !nextImage || (nextImage.regions && nextImage.regions.length > 0) - + const selectedImages = state.images.filter((image) => image.selected) return ( { return new Promise((resolve, reject) => { axios.post(`${config.SERVER_URL}/${api}`, configuration, { responseType: 'blob' }) .then(response => { - const url = window.URL.createObjectURL(new Blob([response.data])); + const url = window.URL.createObjectURL(new Blob([response.data], { type: response.data.type })); resolve(url); }) .catch(error => { diff --git a/client/src/workspace/DownloadButton/DownloadButton.test.js b/client/src/workspace/DownloadButton/DownloadButton.test.js index d220faf..ca8a350 100644 --- a/client/src/workspace/DownloadButton/DownloadButton.test.js +++ b/client/src/workspace/DownloadButton/DownloadButton.test.js @@ -32,7 +32,7 @@ describe('DownloadButton', () => { const mockHandleDownload = jest.fn(); const selectedImageName = 'example.png'; const classList = ['class1', 'class2', 'class3']; // Example class list - + const selectedImages = [{ src: 'example.png' }] const getImageFileSpy = jest.spyOn(require('../../utils/get-data-from-server.js'), 'getImageFile'); const { getByText } = render( @@ -41,6 +41,7 @@ describe('DownloadButton', () => { classList={classList} hideHeaderText={false} disabled={false} + selectedImages= {selectedImages} handleDownload={mockHandleDownload} /> ); diff --git a/client/src/workspace/DownloadButton/index.jsx b/client/src/workspace/DownloadButton/index.jsx index fe117a9..c1af14b 100644 --- a/client/src/workspace/DownloadButton/index.jsx +++ b/client/src/workspace/DownloadButton/index.jsx @@ -14,8 +14,11 @@ import HeaderButton from "../HeaderButton/index.jsx"; import { useTranslation } from "react-i18next" import config from "../../config.js"; -const DownloadButton = ({selectedImageName, classList, hideHeaderText, disabled}) => { +const DownloadButton = ({selectedImageName, classList, hideHeaderText, disabled, selectedImages}) => { const [anchorEl, setAnchorEl] = useState(null); + const getImageNames = () => { + return selectedImages.map(image => image.src.split('/').pop()) + } const { showSnackbar } = useSnackbar(); const {t} = useTranslation(); const handleClick = (event) => { @@ -35,6 +38,7 @@ const DownloadButton = ({selectedImageName, classList, hideHeaderText, disabled} const handleDownload = (format) => { const config_data = {} config_data['image_name'] = selectedImageName + config_data['image_names'] = getImageNames() config_data['colorMap'] = classColorMap config_data['outlineThickness'] = config.OUTLINE_THICKNESS_CONFIG let url = "" @@ -64,6 +68,10 @@ const DownloadButton = ({selectedImageName, classList, hideHeaderText, disabled} link.setAttribute('download', `config_${withoutExtension}.json`); } else if (format === "yolo-annotations") { link.setAttribute('download', `yolo_${withoutExtension}.txt`); + } else if (format === "masked-image") { + link.setAttribute('download', `${withoutExtension}_mask`); + } else if (format == "annotated-image"){ + link.setAttribute('download', `${withoutExtension}_annotated`); } else { link.setAttribute('download', `${withoutExtension}_${format}.png`); } diff --git a/client/src/workspace/Header/Header.test.js b/client/src/workspace/Header/Header.test.js index a75373e..2ff046c 100644 --- a/client/src/workspace/Header/Header.test.js +++ b/client/src/workspace/Header/Header.test.js @@ -31,8 +31,7 @@ describe("Header", () => { ]; const mockClass = ['class1', 'class2', 'class3'] - - + const selectedImages = [] beforeEach(() => { render(
{ onClickItem={mockOnClickItem} selectedImageName="image1" classList={mockClass} + selectedImages={selectedImages} /> ); }); @@ -79,6 +79,7 @@ describe("Header", () => { onClickItem={mockOnClickItem} selectedImageName="image1" classList={mockClass} + selectedImages={selectedImages} /> ); diff --git a/client/src/workspace/Header/index.jsx b/client/src/workspace/Header/index.jsx index 2a4a101..2444463 100644 --- a/client/src/workspace/Header/index.jsx +++ b/client/src/workspace/Header/index.jsx @@ -29,12 +29,14 @@ export const Header = ({ items, onClickItem, selectedImageName, - classList + classList, + selectedImages }) => { const{ t } = useTranslation() const downloadMenu = items.find((item) => item.name === "Download") - + const isDownloadDisabled= downloadMenu && downloadMenu.disabled && selectedImages && selectedImages.length <= 0 + return ( @@ -43,7 +45,7 @@ export const Header = ({ {leftSideContent} {downloadMenu && } {items.filter(item => item.name !== "Download").map((item) => ( diff --git a/client/src/workspace/Workspace/index.jsx b/client/src/workspace/Workspace/index.jsx index 7d864b1..ec9a088 100644 --- a/client/src/workspace/Workspace/index.jsx +++ b/client/src/workspace/Workspace/index.jsx @@ -42,6 +42,7 @@ export default ({ hideHeader = false, hideHeaderText = false, children, + selectedImages = emptyAr, selectedImageName, classList }) => { @@ -58,6 +59,7 @@ export default ({ items={headerItems} selectedImageName={selectedImageName} classList={classList} + selectedImages={selectedImages} /> )} diff --git a/server/app.py b/server/app.py index b3dbed7..9a8a05f 100644 --- a/server/app.py +++ b/server/app.py @@ -1,3 +1,4 @@ +import io from flask import Flask, jsonify, request, url_for,send_from_directory,send_file from flask_cors import CORS, cross_origin from db.db_handler import Module @@ -10,6 +11,7 @@ import traceback import tempfile import shutil +import zipfile app = Flask(__name__) app.config.from_object("config") @@ -244,15 +246,28 @@ def download_configuration(): try: data = request.get_json() # Ensure the expected structure of the JSON data - image_name = data.get('image_name') color_map = data.get("colorMap", None) + image_names = data.get('image_names', []) - if not image_name: - raise ValueError("Invalid JSON data format: 'image_name' not found.") - - json_bytes, download_filename = create_json_response(image_name, color_map) - - return send_file(json_bytes, mimetype='application/json', as_attachment=True, download_name=download_filename) + # Iterate through each image name + all_data = {} + if not image_names: + raise ValueError("Invalid JSON data format: 'image_names' not found.") + + for img_name in image_names: + json_bytes, download_filename = create_json_response(img_name, color_map) + json_data = json_bytes.getvalue().decode('utf-8') + all_data[img_name] = json.loads(json_data) + + # Convert accumulated data to JSON string + json_str = json.dumps(all_data, indent=4) + + return send_file( + io.BytesIO(json_str.encode('utf-8')), + mimetype='application/json', + as_attachment=True, + download_name='merged_configuration.json' + ) except Exception as e: print('Error:', e) @@ -261,77 +276,113 @@ def download_configuration(): @app.route('/download_image_with_annotations', methods=['POST']) @cross_origin(origin=client_url, headers=['Content-Type']) def download_image_with_annotations(): + temp_dir = None # Initialize temporary directory variable try: data = request.get_json() # Ensure the expected structure of the JSON data - image_name = data.get('image_name') - if not image_name: - raise ValueError("Invalid JSON data format: 'image_name' not found.") + image_names = data.get('image_names', []) + if not image_names: + raise ValueError("Invalid JSON data format: 'image_names' not found.") - json_bytes, download_filename = create_json_response(image_name) - # Convert BytesIO to string and return as JSON - json_str = json_bytes.getvalue().decode('utf-8') - - images = json.loads(json_str).get("configuration", []) + # Prepare to store all processed images + img_byte_arrs = [] - color_map = data.get("colorMap", {}) - outlineThickness = data.get("outlineThickness", {}) - - # Convert color map values to tuples - for key in color_map.keys(): - color_map[key] = tuple(color_map[key]) - - for image_info in images: - image_url = image_info.get("regions", [])[0].get("image-src") - - # Docker container uses port 5000, so replace 5001 with 5000 - if "127.0.0.1:5001" in image_url: - image_url = image_url.replace("127.0.0.1:5001", "127.0.0.1:5000") - - response = requests.get(image_url) - image = Image.open(BytesIO(response.content)) - draw = ImageDraw.Draw(image) - - for region in image_info.get("regions", []): - points = region.get("points", []) - width, height = image.size - label = region.get("class") - color = color_map.get(label, (255, 0, 0)) # Default to red if label not in color_map - if 'points' in region and region['points']: - points = region['points'] - scaled_points = [(x * width, y * height) for x, y in points] - # Draw polygon with thicker outline - draw.line(scaled_points + [scaled_points[0]], fill=color, width=outlineThickness.get('POLYGON', 2)) # Change width as desired - elif all(key in region for key in ('x', 'y', 'w', 'h')): - try: - x = float(region['x'][1:-1]) * width if isinstance(region['x'], str) else float(region['x'][0]) * width - y = float(region['y'][1:-1]) * height if isinstance(region['y'], str) else float(region['y'][0]) * height - w = float(region['w'][1:-1]) * width if isinstance(region['w'], str) else float(region['w'][0]) * width - h = float(region['h'][1:-1]) * height if isinstance(region['h'], str) else float(region['h'][0]) * height - except (ValueError, TypeError) as e: - raise ValueError(f"Invalid format in region dimensions: {region}, Error: {e}") - # Draw rectangle with thicker outline - draw.rectangle([x, y, x + w, y + h], outline=color, width=outlineThickness.get('BOUNDING_BOX', 2)) - elif all(key in region for key in ('rx', 'ry', 'rw', 'rh')): - try: - rx = float(region['rx'][1:-1]) * width if isinstance(region['rx'], str) else float(region['rx'][0]) * width - ry = float(region['ry'][1:-1]) * height if isinstance(region['ry'], str) else float(region['ry'][0]) * height - rw = float(region['rw'][1:-1]) * width if isinstance(region['rw'], str) else float(region['rw'][0]) * width - rh = float(region['rh'][1:-1]) * height if isinstance(region['rh'], str) else float(region['rh'][0]) * height - except (ValueError, TypeError) as e: - raise ValueError(f"Invalid format in region dimensions: {region}, Error: {e}") - # Draw ellipse (circle if rw and rh are equal) - draw.ellipse([rx, ry, rx + rw, ry + rh], outline=color, width=outlineThickness.get('CIRCLE', 2)) - - - - img_byte_arr = BytesIO() - image.save(img_byte_arr, format='PNG') - img_byte_arr.seek(0) + for image_name in image_names: + json_bytes, download_filename = create_json_response(image_name) + # Convert BytesIO to string and return as JSON + json_str = json_bytes.getvalue().decode('utf-8') + + images = json.loads(json_str).get("configuration", []) + + color_map = data.get("colorMap", {}) + outlineThickness = data.get("outlineThickness", {}) + + # Convert color map values to tuples + for key in color_map.keys(): + color_map[key] = tuple(color_map[key]) + + for image_info in images: + # image_url = image_info.get("regions", [])[0].get("image-src") + regions = image_info.get("regions", []) + if not regions: + continue # Skip if no regions are present + + region = regions[0] # Take the first region (assuming there is at least one) + image_url = region.get("image-src") + # Docker container uses port 5000, so replace 5001 with 5000 + if "127.0.0.1:5001" in image_url: + image_url = image_url.replace("127.0.0.1:5001", "127.0.0.1:5000") + + response = requests.get(image_url) + image = Image.open(BytesIO(response.content)) + draw = ImageDraw.Draw(image) + + for region in image_info.get("regions", []): + points = region.get("points", []) + width, height = image.size + label = region.get("class") + color = color_map.get(label, (255, 0, 0)) # Default to red if label not in color_map + if 'points' in region and region['points']: + points = region['points'] + scaled_points = [(x * width, y * height) for x, y in points] + # Draw polygon with thicker outline + draw.line(scaled_points + [scaled_points[0]], fill=color, + width=outlineThickness.get('POLYGON', 2)) # Change width as desired + elif all(key in region for key in ('x', 'y', 'w', 'h')): + try: + x = float(region['x'][1:-1]) * width if isinstance(region['x'], str) else float( + region['x'][0]) * width + y = float(region['y'][1:-1]) * height if isinstance(region['y'], str) else float( + region['y'][0]) * height + w = float(region['w'][1:-1]) * width if isinstance(region['w'], str) else float( + region['w'][0]) * width + h = float(region['h'][1:-1]) * height if isinstance(region['h'], str) else float( + region['h'][0]) * height + except (ValueError, TypeError) as e: + raise ValueError(f"Invalid format in region dimensions: {region}, Error: {e}") + # Draw rectangle with thicker outline + draw.rectangle([x, y, x + w, y + h], outline=color, + width=outlineThickness.get('BOUNDING_BOX', 2)) + elif all(key in region for key in ('rx', 'ry', 'rw', 'rh')): + try: + rx = float(region['rx'][1:-1]) * width if isinstance(region['rx'], str) else float( + region['rx'][0]) * width + ry = float(region['ry'][1:-1]) * height if isinstance(region['ry'], str) else float( + region['ry'][0]) * height + rw = float(region['rw'][1:-1]) * width if isinstance(region['rw'], str) else float( + region['rw'][0]) * width + rh = float(region['rh'][1:-1]) * height if isinstance(region['rh'], str) else float( + region['rh'][0]) * height + except (ValueError, TypeError) as e: + raise ValueError(f"Invalid format in region dimensions: {region}, Error: {e}") + # Draw ellipse (circle if rw and rh are equal) + draw.ellipse([rx, ry, rx + rw, ry + rh], outline=color, + width=outlineThickness.get('CIRCLE', 2)) + + img_byte_arr = BytesIO() + image.save(img_byte_arr, format='PNG') + img_byte_arr.seek(0) + + img_byte_arrs.append({ + 'data': img_byte_arr, + 'download_name': image_info.get("image-name") + }) + + # Prepare and return the zip file containing all processed images + zip_byte_arr = BytesIO() + with zipfile.ZipFile(zip_byte_arr, 'w') as zip_file: + for img_info in img_byte_arrs: + zip_file.writestr(f"{img_info['download_name']}.png", img_info['data'].read()) + + zip_byte_arr.seek(0) + + # Clean up temporary resources after sending the zip file + for img_info in img_byte_arrs: + img_info['data'].close() + + return send_file(zip_byte_arr, mimetype='application/zip', as_attachment=True, download_name='images_with_annotations.zip') - return send_file(img_byte_arr, mimetype='image/png', as_attachment=True, download_name=image_info.get("image-name")) - except ValueError as ve: print('ValueError:', ve) traceback.print_exc() @@ -344,6 +395,17 @@ def download_image_with_annotations(): print('General error:', e) traceback.print_exc() return jsonify({'error': str(e)}), 500 + finally: + # Clean up temporary directory if it was created + if temp_dir and os.path.exists(temp_dir): + try: + for file in os.listdir(temp_dir): + file_path = os.path.join(temp_dir, file) + if os.path.isfile(file_path): + os.remove(file_path) + os.rmdir(temp_dir) + except Exception as cleanup_error: + print(f"Error cleaning up temporary directory: {cleanup_error}") @app.route('/download_image_mask', methods=['POST']) @@ -352,69 +414,85 @@ def download_image_mask(): try: data = request.get_json() # Ensure the expected structure of the JSON data - image_name = data.get('image_name') - if not image_name: - raise ValueError("Invalid JSON data format: 'image_name' not found.") - - json_bytes, download_filename = create_json_response(image_name) - # Convert BytesIO to string and return as JSON - json_str = json_bytes.getvalue().decode('utf-8') - - images = json.loads(json_str).get("configuration", []) - - color_map = data.get("colorMap", {}) - outlineThickness = data.get("outlineThickness", {}) - - # Convert color map values to tuples - for key in color_map.keys(): - color_map[key] = tuple(color_map[key]) - - for image_info in images: - image_url = image_info.get("regions", [])[0].get("image-src") - - # Docker container uses port 5000, so replace 5001 with 5000 - if "127.0.0.1:5001" in image_url: - image_url = image_url.replace("127.0.0.1:5001", "127.0.0.1:5000") - - response = requests.get(image_url) - image = Image.open(BytesIO(response.content)) - width, height = image.size - mask = Image.new('RGB', (width, height), app.config["MASK_BACKGROUND_COLOR"]) # 'RGB' mode for colored masks - draw = ImageDraw.Draw(mask) - - for region in image_info.get("regions", []): - label = region.get("class") - color = color_map.get(label, (255, 255, 255)) # Default to white if label not in color_map - if 'points' in region and region['points']: - points = region['points'] - scaled_points = [(int(x * width), int(y * height)) for x, y in points] - draw.polygon(scaled_points, outline=color, fill=color, width=outlineThickness.get('POLYGON', 2)) - elif all(key in region for key in ('x', 'y', 'w', 'h')): - try: - x = float(region['x'][1:-1]) * width if isinstance(region['x'], str) else float(region['x'][0]) * width - y = float(region['y'][1:-1]) * height if isinstance(region['y'], str) else float(region['y'][0]) * height - w = float(region['w'][1:-1]) * width if isinstance(region['w'], str) else float(region['w'][0]) * width - h = float(region['h'][1:-1]) * height if isinstance(region['h'], str) else float(region['h'][0]) * height - except (ValueError, TypeError) as e: - raise ValueError(f"Invalid format in region dimensions: {region}, Error: {e}") - # Draw rectangle for bounding box - draw.rectangle([x, y, x + w, y + h], outline=color, fill=color, width=outlineThickness.get('BOUNDING_BOX', 2)) - elif all(key in region for key in ('rx', 'ry', 'rw', 'rh')): - try: - rx = float(region['rx'][1:-1]) * width if isinstance(region['rx'], str) else float(region['rx'][0]) * width - ry = float(region['ry'][1:-1]) * height if isinstance(region['ry'], str) else float(region['ry'][0]) * height - rw = float(region['rw'][1:-1]) * width if isinstance(region['rw'], str) else float(region['rw'][0]) * width - rh = float(region['rh'][1:-1]) * height if isinstance(region['rh'], str) else float(region['rh'][0]) * height - except (ValueError, TypeError) as e: - raise ValueError(f"Invalid format in region dimensions: {region}, Error: {e}") - # Draw ellipse (circle if rw and rh are equal) - draw.ellipse([rx, ry, rx + rw, ry + rh], outline=color,width=outlineThickness.get('CIRCLE', 2), fill=color) - - mask_byte_arr = BytesIO() - mask.save(mask_byte_arr, format='PNG') - mask_byte_arr.seek(0) - - return send_file(mask_byte_arr, mimetype='image/png', as_attachment=True, download_name=f"mask_{image_info.get('image-name')}") + image_names = data.get('image_names', []) + if not image_names: + raise ValueError("Invalid JSON data format: 'image_names' not found.") + + # Initialize a temporary directory to store mask images + temp_dir = tempfile.mkdtemp() + + # Process each image and create masks + zip_filename = 'image_masks.zip' + zip_file_path = os.path.join(temp_dir, zip_filename) + + with zipfile.ZipFile(zip_file_path, 'w') as zipf: + for image_name in image_names: + json_bytes, download_filename = create_json_response(image_name) + json_str = json_bytes.getvalue().decode('utf-8') + images = json.loads(json_str).get("configuration", []) + + color_map = data.get("colorMap", {}) + outlineThickness = data.get("outlineThickness", {}) + + # Convert color map values to tuples + for key in color_map.keys(): + color_map[key] = tuple(color_map[key]) + for image_info in images: + regions = image_info.get("regions", []) + if not regions: + continue # Skip if no regions are present + + region = regions[0] # Take the first region (assuming there is at least one) + image_url = region.get("image-src") + # Docker container uses port 5000, so replace 5001 with 5000 + if "127.0.0.1:5001" in image_url: + image_url = image_url.replace("127.0.0.1:5001", "127.0.0.1:5000") + + response = requests.get(image_url) + response.raise_for_status() + image = Image.open(BytesIO(response.content)) + width, height = image.size + mask = Image.new('RGB', (width, height), app.config["MASK_BACKGROUND_COLOR"]) # 'RGB' mode for colored masks + draw = ImageDraw.Draw(mask) + + for region in image_info.get("regions", []): + label = region.get("class") + color = color_map.get(label, (255, 255, 255)) # Default to white if label not in color_map + if 'points' in region and region['points']: + points = region['points'] + scaled_points = [(int(x * width), int(y * height)) for x, y in points] + draw.polygon(scaled_points, outline=color, fill=color, width=outlineThickness.get('POLYGON', 2)) + elif all(key in region for key in ('x', 'y', 'w', 'h')): + try: + x = float(region['x'][1:-1]) * width if isinstance(region['x'], str) else float(region['x'][0]) * width + y = float(region['y'][1:-1]) * height if isinstance(region['y'], str) else float(region['y'][0]) * height + w = float(region['w'][1:-1]) * width if isinstance(region['w'], str) else float(region['w'][0]) * width + h = float(region['h'][1:-1]) * height if isinstance(region['h'], str) else float(region['h'][0]) * height + except (ValueError, TypeError) as e: + raise ValueError(f"Invalid format in region dimensions: {region}, Error: {e}") + # Draw rectangle for bounding box + draw.rectangle([x, y, x + w, y + h], outline=color, fill=color, width=outlineThickness.get('BOUNDING_BOX', 2)) + elif all(key in region for key in ('rx', 'ry', 'rw', 'rh')): + try: + rx = float(region['rx'][1:-1]) * width if isinstance(region['rx'], str) else float(region['rx'][0]) * width + ry = float(region['ry'][1:-1]) * height if isinstance(region['ry'], str) else float(region['ry'][0]) * height + rw = float(region['rw'][1:-1]) * width if isinstance(region['rw'], str) else float(region['rw'][0]) * width + rh = float(region['rh'][1:-1]) * height if isinstance(region['rh'], str) else float(region['rh'][0]) * height + except (ValueError, TypeError) as e: + raise ValueError(f"Invalid format in region dimensions: {region}, Error: {e}") + # Draw ellipse (circle if rw and rh are equal) + draw.ellipse([rx, ry, rx + rw, ry + rh], outline=color, width=outlineThickness.get('CIRCLE', 2), fill=color) + + # Save mask image to temporary directory + mask_filename = f"mask_{image_info.get('image-name').split('.')[0]}.png" + mask_path = os.path.join(temp_dir, mask_filename) + mask.save(mask_path, format='PNG') + + # Add mask image to zip file + zipf.write(mask_path, arcname=mask_filename) + + # Send zip file as response + return send_file(zip_file_path, mimetype='application/zip', as_attachment=True, download_name=zip_filename) except ValueError as ve: print('ValueError:', ve) @@ -428,111 +506,126 @@ def download_image_mask(): print('General error:', e) traceback.print_exc() return jsonify({'error': str(e)}), 500 - -def create_yolo_annotations(image_name, color_map=None): + finally: + # Clean up temporary directory + try: + if temp_dir: + for file in os.listdir(temp_dir): + file_path = os.path.join(temp_dir, file) + if os.path.isfile(file_path): + os.remove(file_path) + os.rmdir(temp_dir) + except Exception as e: + print(f"Error cleaning up temporary directory: {e}") + +def create_yolo_annotations(image_names, color_map=None): base_url = request.host_url + 'uploads/' - annotations = [] - - # Fetch image and its annotations - image_path = None - for root, dirs, files in os.walk(path): - for f in files: - if f.lower().endswith(('.png', '.jpg', '.jpeg')) and f.lower() == image_name.lower(): - image_path = os.path.join(root, f) + all_annotations = [] + + for image_name in image_names: + # Fetch image and its annotations + image_path = None + for root, dirs, files in os.walk(path): + for f in files: + if f.lower().endswith(('.png', '.jpg', '.jpeg')) and f.lower() == image_name.lower(): + image_path = os.path.join(root, f) + break + if image_path: break - if image_path: - break - - if not image_path: - raise ValueError(f"Image '{image_name}' not found in the upload directory.") - - # Fetch annotations from database or other source based on image_path - image_url = base_url + image_name - imageIndex = dbModule.findInfoInDb(dbModule.imagesInfo, 'image-src', image_url) - polygonRegions = dbModule.findInfoInPolygonDb(dbModule.imagePolygonRegions, 'image-src', image_url) - boxRegions = dbModule.findInfoInBoxDb(dbModule.imageBoxRegions, 'image-src', image_url) - circleRegions = dbModule.findInfoInCircleDb(dbModule.imageCircleRegions, 'image-src', image_url) - - # Initialize YOLO annotations - width, height = Image.open(image_path).size - annotations = [] - - # Process polygon regions - if polygonRegions is not None: - for index, region in polygonRegions.iterrows(): - class_name = region.get('class', 'unknown') - points_str = region.get('points', '') - - # Split points string into individual points - points_list = points_str.split(';') - - # Convert points to list of tuples - points = [] - for point_str in points_list: - x, y = map(float, point_str.split('-')) - points.append((x, y)) - - # Convert points to normalized YOLO format - if points: - xmin = min(point[0] for point in points) / width - ymin = min(point[1] for point in points) / height - xmax = max(point[0] for point in points) / width - ymax = max(point[1] for point in points) / height + + if not image_path: + raise ValueError(f"Image '{image_name}' not found in the upload directory.") + + # Fetch annotations from database or other source based on image_path + image_url = base_url + image_name + imageIndex = dbModule.findInfoInDb(dbModule.imagesInfo, 'image-src', image_url) + polygonRegions = dbModule.findInfoInPolygonDb(dbModule.imagePolygonRegions, 'image-src', image_url) + boxRegions = dbModule.findInfoInBoxDb(dbModule.imageBoxRegions, 'image-src', image_url) + circleRegions = dbModule.findInfoInCircleDb(dbModule.imageCircleRegions, 'image-src', image_url) + + # Initialize YOLO annotations for current image + width, height = Image.open(image_path).size + annotations = [] + + # Process polygon regions + if polygonRegions is not None: + for index, region in polygonRegions.iterrows(): + class_name = region.get('class', 'unknown') + points_str = region.get('points', '') - # YOLO format: class_index x_center y_center width height (all normalized) - annotations.append(f"{class_name} {(xmin + xmax) / 2} {(ymin + ymax) / 2} {xmax - xmin} {ymax - ymin}") - - # Process box regions - if boxRegions is not None: - for index, region in boxRegions.iterrows(): - class_name = region.get('class', 'unknown') - try: - x = float(region['x'][1:-1]) * width if isinstance(region['x'], str) else float(region['x'][0]) * width - y = float(region['y'][1:-1]) * height if isinstance(region['y'], str) else float(region['y'][0]) * height - w = float(region['w'][1:-1]) * width if isinstance(region['w'], str) else float(region['w'][0]) * width - h = float(region['h'][1:-1]) * height if isinstance(region['h'], str) else float(region['h'][0]) * height - except (ValueError, TypeError) as e: - raise ValueError(f"Invalid format in region dimensions: {region}, Error: {e}") - - # YOLO format: class_index x_center y_center width height (all normalized) - annotations.append(f"{class_name} {x + w / 2} {y + h / 2} {w} {h}") - - # Process circle/ellipse regions - if circleRegions is not None: - for index, region in circleRegions.iterrows(): - class_name = region.get('class', 'unknown') - try: - rx = float(region['rx'][1:-1]) * width if isinstance(region['rx'], str) else float(region['rx'][0]) * width - ry = float(region['ry'][1:-1]) * height if isinstance(region['ry'], str) else float(region['ry'][0]) * height - rw = float(region['rw'][1:-1]) * width if isinstance(region['rw'], str) else float(region['rw'][0]) * width - rh = float(region['rh'][1:-1]) * height if isinstance(region['rh'], str) else float(region['rh'][0]) * height - except (ValueError, TypeError) as e: - raise ValueError(f"Invalid format in region dimensions: {region}, Error: {e}") - - # For YOLO, if width and height are equal, it represents a circle - if rw == rh: - annotations.append(f"{class_name} {rx} {ry} {rw} {rw}") # Treat as circle - else: - # Treat as ellipse (YOLO does not directly support ellipse, so treat as box) - annotations.append(f"{class_name} {rx + rw / 2} {ry + rh / 2} {rw} {rh}") + # Split points string into individual points + points_list = points_str.split(';') + + # Convert points to list of tuples + points = [] + for point_str in points_list: + x, y = map(float, point_str.split('-')) + points.append((x, y)) + + # Convert points to normalized YOLO format + if points: + xmin = min(point[0] for point in points) / width + ymin = min(point[1] for point in points) / height + xmax = max(point[0] for point in points) / width + ymax = max(point[1] for point in points) / height + + # YOLO format: class_index x_center y_center width height (all normalized) + annotations.append(f"{class_name} {(xmin + xmax) / 2} {(ymin + ymax) / 2} {xmax - xmin} {ymax - ymin}") + + # Process box regions + if boxRegions is not None: + for index, region in boxRegions.iterrows(): + class_name = region.get('class', 'unknown') + try: + x = float(region['x'][1:-1]) * width if isinstance(region['x'], str) else float(region['x'][0]) * width + y = float(region['y'][1:-1]) * height if isinstance(region['y'], str) else float(region['y'][0]) * height + w = float(region['w'][1:-1]) * width if isinstance(region['w'], str) else float(region['w'][0]) * width + h = float(region['h'][1:-1]) * height if isinstance(region['h'], str) else float(region['h'][0]) * height + except (ValueError, TypeError) as e: + raise ValueError(f"Invalid format in region dimensions: {region}, Error: {e}") - print(f"Annotations: {annotations}") - return annotations + # YOLO format: class_index x_center y_center width height (all normalized) + annotations.append(f"{class_name} {x + w / 2} {y + h / 2} {w} {h}") + + # Process circle/ellipse regions + if circleRegions is not None: + for index, region in circleRegions.iterrows(): + class_name = region.get('class', 'unknown') + try: + rx = float(region['rx'][1:-1]) * width if isinstance(region['rx'], str) else float(region['rx'][0]) * width + ry = float(region['ry'][1:-1]) * height if isinstance(region['ry'], str) else float(region['ry'][0]) * height + rw = float(region['rw'][1:-1]) * width if isinstance(region['rw'], str) else float(region['rw'][0]) * width + rh = float(region['rh'][1:-1]) * height if isinstance(region['rh'], str) else float(region['rh'][0]) * height + except (ValueError, TypeError) as e: + raise ValueError(f"Invalid format in region dimensions: {region}, Error: {e}") + + # For YOLO, if width and height are equal, it represents a circle + if rw == rh: + annotations.append(f"{class_name} {rx} {ry} {rw} {rw}") # Treat as circle + else: + # Treat as ellipse (YOLO does not directly support ellipse, so treat as box) + annotations.append(f"{class_name} {rx + rw / 2} {ry + rh / 2} {rw} {rh}") + + # Append annotations for current image to all_annotations list + all_annotations.extend(annotations) + + print(f"Annotations: {all_annotations}") + return all_annotations @app.route('/download_yolo_annotations', methods=['POST']) @cross_origin(origin=client_url, headers=['Content-Type']) def download_yolo_annotations(): data = request.get_json() - image_name = data.get('image_name') + image_names = data.get('image_names') - if not image_name: - return jsonify({'error': "Invalid JSON data format: 'image_name' not found."}), 400 + if not image_names: + return jsonify({'error': "Invalid JSON data format: 'image_names' not found."}), 400 temp_file = None try: - annotations = create_yolo_annotations(image_name) + annotations = create_yolo_annotations(image_names) print(f"Annotations: {annotations}") # Create a temporary text file with YOLO annotations @@ -544,7 +637,7 @@ def download_yolo_annotations(): f.write(annotation + "\n") # Set a meaningful download name - download_name = f"{image_name.split('.')[0]}.txt" + download_name = 'yolo_annotations.txt' return send_file(temp_file.name, as_attachment=True, download_name=download_name), 200 diff --git a/server/tests/test_app.py b/server/tests/test_app.py index fc7eb92..9b0019a 100644 --- a/server/tests/test_app.py +++ b/server/tests/test_app.py @@ -86,22 +86,22 @@ def test_images_name_no_image_name(self): def test_download_configuration_no_image_name(self): response = self.app.post('/download_configuration', data=json.dumps({}), content_type='application/json') self.assertEqual(response.status_code, 500) - self.assertIn(b"'image_name' not found", response.data) + self.assertIn(b"'image_names' not found", response.data) def test_download_image_with_annotations_no_image_name(self): response = self.app.post('/download_image_with_annotations', data=json.dumps({}), content_type='application/json') self.assertEqual(response.status_code, 400) - self.assertIn(b"'image_name' not found", response.data) + self.assertIn(b"'image_names' not found", response.data) def test_download_image_mask_no_image_name(self): response = self.app.post('/download_image_mask', data=json.dumps({}), content_type='application/json') self.assertEqual(response.status_code, 400) - self.assertIn(b"'image_name' not found", response.data) + self.assertIn(b"'image_names' not found", response.data) def test_download_yolo_annotations_no_image_name(self): response = self.app.post('/download_yolo_annotations', data=json.dumps({}), content_type='application/json') self.assertEqual(response.status_code, 400) - self.assertIn(b"'image_name' not found", response.data) + self.assertIn(b"'image_names' not found", response.data) # def test_get_images_info_no_path(self): # app.config['UPLOAD_FOLDER'] = '/nonexistent_path'