[
{
"path": "LICENSE",
"content": "MIT License\n\nCopyright (c) 2024 Dr. Sreenivas Bhattiprolu\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n\n--------------------------------------------------------------------------------\n\nCITATION REQUEST:\n\nIf you use this software in your research, please cite it as follows:\n\nBhattiprolu, S. (2024). Image Annotator [Computer software]. \nhttps://github.com/bnsreenu/digitalsreeni-image-annotator\n\nBibTeX:\n@software{image_annotator,\n author = {Bhattiprolu, Sreenivas},\n title = {Image Annotator},\n year = {2024},\n url = {https://github.com/bnsreenu/digitalsreeni-image-annotator}\n}\n\nWhile not required by the license, citation is appreciated and helps support the \ncontinued development and maintenance of this software."
},
{
"path": "README.md",
"content": "# DigitalSreeni Image Annotator and Toolkit\n\n\n\n\n\nA powerful and user-friendly tool for annotating images with polygons and rectangles, built with PyQt5. Now with additional supporting tools for comprehensive image processing and dataset management.\n\n## Support the Project\n\nIf you find this project helpful, consider supporting it:\n\n[](https://www.paypal.com/donate/?business=FGQL3CNJGJP9C&no_recurring=0&item_name=If+you+find+this+Image+Annotator+project+helpful%2C+consider+supporting+it%3A¤cy_code=USD)\n\n\n\n## Watch the demo (of v0.8.0):\n[](https://youtu.be/aArn1f1YIQk)\n\n@DigitalSreeni\nDr. Sreenivas Bhattiprolu\n\n## Features\n\n- Semi-automated annotations with SAM-2 assistance (Segment Anything Model) — Because who doesn't love a helpful AI sidekick?\n- Manual annotations with polygons and rectangles — For when you want to show SAM-2 who's really in charge.\n- Paint brush and Eraser tools with adjustable pen sizes (use - and = on your keyboard)\n- Merge annotations - For when SAM-2's guesswork needs a little human touch. \n- Save and load projects for continued work.\n- Save As... and Autosave functionality. \n- A secret game, for when you are bored.\n- Import existing COCO JSON annotations with images.\n- Export annotations to various formats (COCO JSON, YOLO v8/v11, Labeled images, Semantic labels, Pascal VOC).\n- Handle multi-dimensional images (TIFF stacks and CZI files).\n- Zoom and pan for detailed annotations.\n- Support for multiple classes with customizable colors.\n- User-friendly interface with intuitive controls.\n- Change the application font size on the fly — Make your annotations as big or small as your caffeine level requires.\n- Dark mode for those late-night annotation marathons — Who needs sleep when you have dark mode?\n- Pick appropriate pre-trained SAM2 model for flexible and improved semi-automated annotations.\n- Change the class of an annotation to a different class.\n- Turn visibility of a class ON and OFF.\n- YOLO (beta) training using current annotations and loading trained model to segment images.\n- Area measurements for annotations displayed next to the Annotation name.\n- Sort annotations by name/number or area.\n- Additional supporting tools:\n - Annotation statistics for current annotations\n - COCO JSON combiner\n - Dataset splitter\n - Stack to slices converter\n - Image patcher\n - Image augmenter\n- Project Details: View and edit project metadata, including creation date, last modified date, image information, and custom notes.\n- Advanced Project Search: Search through multiple projects using complex queries with logical operators (AND, OR) and parentheses.\n- Slice Registration\n - Align image slices in a stack with multiple registration methods\n - Support for various reference frames and transformation types\n - Stack Interpolation\n - Adjust Z-spacing in image stacks\n - Multiple interpolation methods with memory-efficient processing\n - DICOM Converter\n - Convert DICOM files to TIFF format (single stack or individual slices)\n - Preserve metadata and physical dimensions\n - Export metadata to JSON for reference\n\n\n## Operating System Requirements\nThis application is built using PyQt5 and has been tested on macOS and Windows. It may experience compatibility issues on Linux systems, particularly related to the XCB plugin for PyQt5. Extensive testing on Linux systems has not been done yet.\n\n## Installation\n\n### Watch the installation walkthough video:\n[](https://youtu.be/VI6V95eUUpY)\n\nYou can install the DigitalSreeni Image Annotator directly from PyPI:\n\n```bash\npip install digitalsreeni-image-annotator\n```\n\nThe application uses the Ultralytics library, so there's no need to separately install SAM2 or PyTorch, or download SAM2 models manually.\n\n## Usage\n\n1. Run the DigitalSreeni Image Annotator application:\n ```bash\n digitalsreeni-image-annotator\n ```\n or\n ```bash\n sreeni\n ```\n or\n ```bash\n python -m digitalsreeni_image_annotator.main\n ```\n\n2. Using the application:\n - Click \"New Project\" or use Ctrl+N to start a new project.\n - Use \"Add New Images\" to import images, including TIFF stacks and CZI files.\n - Add classes using the \"Add Classes\" button.\n - Select a class and use the Polygon or Rectangle or Paint Brush tool to create manual annotations.\n - To use SAM2-assisted annotation:\n - Select a model from the \"Pick a SAM Model\" dropdown. It's recommended to use smaller models like SAM2 tiny or SAM2 small. SAM2 large is not recommended as it may crash the application on systems with limited resources. \n - Note: When you select a model for the first time, the application needs to download it. This process may take a few seconds to a minute, depending on your internet connection speed. Subsequent uses of the same model will be faster as it will already be cached locally, in your working directory.\n - Click the \"SAM-Assisted\" button to activate the tool.\n - Draw a rectangle around objects of interest to allow SAM2 to automatically detect objects.\n - Note that SAM2 provides various outputs with different scores, and only the top-scoring region will be displayed. If the desired result isn't achieved on the first try, draw again.\n - For low-quality images where SAM2 may not auto-detect objects, manual tools may be necessary.\n - When SAM2 auto-detect partial objects, use polygon or paint brush tools to manually define the remaining region and use the Merge tool to combine both annotations into one.\n - When SAM2 over-annotates objects, extending the annotation beyond object's boundaries, use the Eraser tool to clean up the edges. \n - Both paint brush and eraser tools can be adjusted for pen size by using - or = keys on your keyboard. \n - Edit existing annotations by double-clicking on them.\n - Edit existing annotations using the Eraser tool. Adjust the eraser size by using - or = keys on your keyboard.\n - Merge connected annotations by selecting them from the Annotations list and clicking the Merge button. \n - Change the class of an annotation to a different class.\n - Turn visibility of a class ON and OFF.\n - Use YOLO (beta) training with current annotations and load the trained model to segment images and convert segmentations to annotations. (Currently not implemented for slices or stacks, just single images.)\n - Accept/reject one or select class predictions at a time to add them as annotations.\n - View area measurements for annotations displayed next to the Annotation name.\n - Sort annotations by name/number or area.\n - Save your project using \"Save Project\" or Ctrl+S. Alternatively, you can use Save As... to save the project with a different name. \n - Use \"Open Project\" or Ctrl+O to load a previously saved project.\n - Click \"Import Annotations with Images\" to load existing COCO JSON annotations along with their images.\n - Use \"Export Annotations\" to save annotations in various formats (COCO JSON, YOLO v8/v11, Labeled images, Semantic labels, Pascal VOC).\n - Note: YOLO export (and import) is now compatible with YOLOv11 structure. (Project directory includes data.yaml, train, and valid directories, with train and valid both having images and labels subdirectories.)\n - Project Details:\n - Access project details by selecting \"Project Details\" from the Project menu.\n - View project metadata such as creation date, last modified date, and image information.\n - Add or edit custom project notes.\n - Project details are automatically saved when you make changes to the notes.\n - Advanced Project Search:\n - Access the search functionality by selecting \"Search Projects\" from the Project menu.\n - Search through multiple projects using complex queries.\n - Use logical operators (AND, OR) and parentheses for advanced search criteria.\n - Search covers project name, class names, image names, and project notes.\n - Example queries:\n - \"cells AND dog\": Find projects containing both \"cells\" and \"dog\"\n - \"cells OR bacteria\": Find projects containing either \"cells\" or \"bacteria\"\n - \"cells AND (dog OR monkey)\": Find projects containing \"cells\" and either \"dog\" or \"monkey\"\n - \"(project1 OR project2) AND (cells OR bacteria)\": More complex nested queries\n - Double-click on search results to open the corresponding project.\n - Access additional tools under the Tools menu bar:\n - Annotation Statistics\n - COCO JSON Combiner\n - Dataset Splitter\n - Stack to Slices Converter\n - Image Patcher\n - Image Augmenter\n - Each tool opens a separate UI to guide you through the respective task.\n - Access the help documentation by clicking the \"Help\" button or pressing F1.\n - Explore the interface – you might stumble upon some hidden gems and secret features!\n\n3. Keyboard shortcuts:\n - Ctrl + N: Create a new project\n - Ctrl + O: Open an existing project\n - Ctrl + S: Save the current project\n - Ctrl + W: Close the current project\n - Ctrl + Shift + S: Open Annotation Statistics\n - F1: Open the help window\n - Ctrl + Wheel: Zoom in/out\n - Hold Ctrl and drag: Pan the image\n - Esc: Cancel current annotation, exit edit mode, or exit SAM-assisted annotation\n - Enter: Finish current annotation, exit edit mode, or accept SAM-generated mask\n - Up/Down Arrow Keys: Navigate through slices in multi-dimensional images\n - - and =: Adjust pen size for paint brush and eraser tools\n\n## Known Issues and Bug Fixes\n\n- The application may not work correctly on Linux systems. Extensive testing has not been done yet.\n- When loading a YOLO model trained on different classes compared to the loaded YAML file, the application now gives a message to the user about the mismatch instead of crashing.\n- Various other bugs have been addressed to improve overall stability and performance.\n\n## Development\n\nFor development purposes, you can clone the repository and install it in editable mode:\n\n1. Clone the repository:\n ```bash\n git clone https://github.com/bnsreenu/digitalsreeni-image-annotator.git\n cd digitalsreeni-image-annotator\n ```\n\n2. Create a virtual environment (optional but recommended):\n ```bash\n python -m venv venv\n source venv/bin/activate # On Windows, use `venv\\Scripts\\activate`\n ```\n\n3. Install the package and its dependencies in editable mode:\n ```bash\n pip install -e .\n ```\n\n## Contributing\n\nContributions are welcome! Please feel free to submit a Pull Request.\n\n1. Fork the repository\n2. Create your feature branch (`git checkout -b feature/AmazingFeature`)\n3. Commit your changes (`git commit -m 'Add some AmazingFeature'`)\n4. Push to the branch (`git push origin feature/AmazingFeature`)\n5. Open a Pull Request\n\n## License\n\nThis project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.\n\n## Acknowledgments\n\n- Thanks to all my [YouTube](http://www.youtube.com/c/DigitalSreeni) subscribers who inspired me to work on this project\n- Inspired by the need for efficient image annotation in computer vision tasks\n\n## Contact\n\nDr. Sreenivas Bhattiprolu - [@DigitalSreeni](https://twitter.com/DigitalSreeni)\n\nProject Link: [https://github.com/bnsreenu/digitalsreeni-image-annotator](https://github.com/bnsreenu/digitalsreeni-image-annotator)\n\n## Citing\n\nIf you use this software in your research, please cite it as follows:\n\nBhattiprolu, S. (2024). DigitalSreeni Image Annotator [Computer software]. \nhttps://github.com/bnsreenu/digitalsreeni-image-annotator\n\n```bibtex\n@software{digitalsreeni_image_annotator,\n author = {Bhattiprolu, Sreenivas},\n title = {DigitalSreeni Image Annotator},\n year = {2024},\n url = {https://github.com/bnsreenu/digitalsreeni-image-annotator}\n}\n```\n"
},
{
"path": "Release Notes 0.8.12.md",
"content": "# Release Notes\n## Version 0.8.12\n\n### New Features and Enhancements\n- Same as version 0.8.9 except changed the requirements file to define specific version numbers for the libraies used.\n- The following bug fixes and optimizations correspond to version 0.8.9 \n\n### Bug Fixes and Optimizations\n1. **Project Corruption Prevention**\n - Fixed critical issue where projects could become corrupted if application was terminated during loading\n - Disabled auto-save functionality during project loading process\n - Enhanced project loading stability for large datasets\n - Protected project integrity when handling multiple classes and images\n\n### Notes\n- All existing tools continue to support both Windows and macOS operating systems\n- Improved reliability of project file handling\n- Critical update recommended for users working with large projects"
},
{
"path": "data/YOLO11n-model-yaml/coco8.yaml",
"content": "\n# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]\npath: ../datasets/coco8 # dataset root dir\ntrain: images/train # train images (relative to 'path') 4 images\nval: images/val # val images (relative to 'path') 4 images\ntest: # test images (optional)\n\n# Classes\nnames:\n 0: person\n 1: bicycle\n 2: car\n 3: motorcycle\n 4: airplane\n 5: bus\n 6: train\n 7: truck\n 8: boat\n 9: traffic light\n 10: fire hydrant\n 11: stop sign\n 12: parking meter\n 13: bench\n 14: bird\n 15: cat\n 16: dog\n 17: horse\n 18: sheep\n 19: cow\n 20: elephant\n 21: bear\n 22: zebra\n 23: giraffe\n 24: backpack\n 25: umbrella\n 26: handbag\n 27: tie\n 28: suitcase\n 29: frisbee\n 30: skis\n 31: snowboard\n 32: sports ball\n 33: kite\n 34: baseball bat\n 35: baseball glove\n 36: skateboard\n 37: surfboard\n 38: tennis racket\n 39: bottle\n 40: wine glass\n 41: cup\n 42: fork\n 43: knife\n 44: spoon\n 45: bowl\n 46: banana\n 47: apple\n 48: sandwich\n 49: orange\n 50: broccoli\n 51: carrot\n 52: hot dog\n 53: pizza\n 54: donut\n 55: cake\n 56: chair\n 57: couch\n 58: potted plant\n 59: bed\n 60: dining table\n 61: toilet\n 62: tv\n 63: laptop\n 64: mouse\n 65: remote\n 66: keyboard\n 67: cell phone\n 68: microwave\n 69: oven\n 70: toaster\n 71: sink\n 72: refrigerator\n 73: book\n 74: clock\n 75: vase\n 76: scissors\n 77: teddy bear\n 78: hair drier\n 79: toothbrush\n\n# Download script/URL (optional)\ndownload: https://github.com/ultralytics/assets/releases/download/v0.0.0/coco8.zip\n"
},
{
"path": "data/YOLO11n-model-yaml/download_YOLO_models.txt",
"content": "https://docs.ultralytics.com/tasks/segment/#models\n\nRecommended model: https://github.com/ultralytics/assets/releases/download/v8.3.0/yolo11n-seg.pt"
},
{
"path": "data/download_SAM_models.txt",
"content": "It is recommended to pre-download SAM models and place them in your working director - the directory from where you are starting this application. This avoids downloading the models multiple times. \n\nDownload models from: https://docs.ultralytics.com/models/sam-2/\n\nDirect Download links:\nTiny model: https://github.com/ultralytics/assets/releases/download/v8.2.0/sam2_t.pt\nSmall: https://github.com/ultralytics/assets/releases/download/v8.2.0/sam2_s.pt\nBase: https://github.com/ultralytics/assets/releases/download/v8.2.0/sam2_b.pt\nLarge: https://github.com/ultralytics/assets/releases/download/v8.2.0/sam2_l.pt\n\nBe cautious with the large model as it demands higher computing and memory resources from your system."
},
{
"path": "requirements.txt",
"content": "PyQt5==5.15.11\nPillow==11.0.0\nnumpy==2.1.3\ntifffile==2023.3.15\nczifile==2019.7.2\nopencv-python==4.10.0.84\npyyaml==6.0.2\nscikit-image==0.24.0\nultralytics==8.3.27\nplotly==5.24.1\nshapely==2.0.6\npystackreg==0.2.8\npydicom==3.0.1"
},
{
"path": "setup.py",
"content": "\"\"\"\nSetup file for the DigitalSreeni Image Annotator package.\n@DigitalSreeni\nDr. Sreenivas Bhattiprolu\n\"\"\"\nfrom setuptools import setup, find_packages\n\nwith open(\"README.md\", \"r\", encoding=\"utf-8\") as fh:\n long_description = fh.read()\n\nsetup(\n name=\"digitalsreeni-image-annotator\",\n version=\"0.8.12\", # Updated version number\n author=\"Dr. Sreenivas Bhattiprolu\",\n author_email=\"digitalsreeni@gmail.com\",\n description=\"A tool for annotating images using manual and automated tools, supporting multi-dimensional images and SAM2-assisted annotations\",\n long_description=long_description,\n long_description_content_type=\"text/markdown\",\n url=\"https://github.com/bnsreenu/digitalsreeni-image-annotator\",\n packages=find_packages(where=\"src\"),\n package_dir={\"\": \"src\"},\n classifiers=[\n \"Development Status :: 3 - Alpha\",\n \"Intended Audience :: Science/Research\",\n \"Intended Audience :: Developers\",\n \"License :: OSI Approved :: MIT License\",\n \"Operating System :: OS Independent\",\n \"Programming Language :: Python :: 3.10\",\n ],\n python_requires=\">=3.10\",\n install_requires=[\n \"PyQt5==5.15.11\",\n \"numpy==2.1.3\",\n \"Pillow==11.0.0\",\n \"tifffile==2023.3.15\",\n \"czifile==2019.7.2\",\n \"opencv-python==4.10.0.84\",\n \"pyyaml==6.0.2\",\n \"scikit-image==0.24.0\",\n \"ultralytics==8.3.27\",\n \"plotly==5.24.1\",\n \"shapely==2.0.6\", \n \"pystackreg==0.2.8\",\n \"pydicom==3.0.1\"\n ],\n entry_points={\n \"console_scripts\": [\n \"digitalsreeni-image-annotator=digitalsreeni_image_annotator.main:main\",\n \"sreeni=digitalsreeni_image_annotator.main:main\", \n ],\n },\n)"
},
{
"path": "src/digitalsreeni_image_annotator/__init__.py",
"content": "\"\"\"\nImage Annotator\n===============\nA tool for annotating images with polygons and rectangles.\nThis package provides a GUI application for image annotation,\nsupporting polygon and rectangle annotations in a COCO-compatible format.\n@DigitalSreeni\nDr. Sreenivas Bhattiprolu\n\"\"\"\n__version__ = \"0.8.12\"\n__author__ = \"Dr. Sreenivas Bhattiprolu\"\n\nfrom .annotator_window import ImageAnnotator\nfrom .image_label import ImageLabel\nfrom .utils import calculate_area, calculate_bbox\nfrom .sam_utils import SAMUtils \n\n__all__ = ['ImageAnnotator', 'ImageLabel', 'calculate_area', 'calculate_bbox', 'SAMUtils'] # Add 'SAMUtils' to this list"
},
{
"path": "src/digitalsreeni_image_annotator/annotation_statistics.py",
"content": "import plotly.graph_objects as go\nfrom plotly.subplots import make_subplots\nfrom PyQt5.QtWidgets import QDialog, QVBoxLayout, QTextBrowser, QPushButton, QHBoxLayout\nfrom PyQt5.QtCore import Qt\nimport tempfile\nimport os\nimport webbrowser\n\nclass AnnotationStatisticsDialog(QDialog):\n def __init__(self, parent=None):\n super().__init__(parent)\n self.setWindowTitle(\"Annotation Statistics\")\n self.setGeometry(100, 100, 600, 400)\n self.setWindowFlags(self.windowFlags() | Qt.Window)\n self.initUI()\n\n def initUI(self):\n layout = QVBoxLayout()\n self.text_browser = QTextBrowser()\n layout.addWidget(self.text_browser)\n\n button_layout = QHBoxLayout()\n self.show_plot_button = QPushButton(\"Show Interactive Plot\")\n self.show_plot_button.clicked.connect(self.show_interactive_plot)\n button_layout.addWidget(self.show_plot_button)\n\n layout.addLayout(button_layout)\n self.setLayout(layout)\n\n self.plot_file = None\n\n def show_centered(self, parent):\n parent_geo = parent.geometry()\n self.move(parent_geo.center() - self.rect().center())\n self.show()\n\n def generate_statistics(self, annotations):\n try:\n # Class distribution\n class_distribution = {}\n objects_per_image = {}\n total_objects = 0\n \n for image, image_annotations in annotations.items():\n objects_in_image = 0\n for class_name, class_annotations in image_annotations.items():\n class_count = len(class_annotations)\n class_distribution[class_name] = class_distribution.get(class_name, 0) + class_count\n objects_in_image += class_count\n total_objects += class_count\n objects_per_image[image] = objects_in_image\n \n avg_objects_per_image = total_objects / len(annotations) if annotations else 0\n \n # Create plots\n fig = make_subplots(rows=2, cols=1, subplot_titles=(\"Class Distribution\", \"Objects per Image\"))\n \n # Class distribution plot\n fig.add_trace(go.Bar(x=list(class_distribution.keys()), y=list(class_distribution.values()), name=\"Classes\"),\n row=1, col=1)\n \n # Objects per image plot\n fig.add_trace(go.Bar(\n x=list(objects_per_image.keys()),\n y=list(objects_per_image.values()),\n name=\"Images\",\n hovertext=[f\"{img}: {count}\" for img, count in objects_per_image.items()],\n hoverinfo=\"text\"\n ), row=2, col=1)\n \n # Update layout\n fig.update_layout(height=800, title_text=\"Annotation Statistics\")\n \n # Hide x-axis labels for the second subplot (Objects per Image)\n fig.update_xaxes(showticklabels=False, title_text=\"Images\", row=2, col=1)\n \n # Update y-axis title for the second subplot\n fig.update_yaxes(title_text=\"Number of Objects\", row=2, col=1)\n \n # Save the plot to a temporary HTML file\n with tempfile.NamedTemporaryFile(mode=\"w\", suffix=\".html\", delete=False) as tmp:\n fig.write_html(tmp.name)\n self.plot_file = tmp.name\n \n # Display statistics in the text browser\n stats_text = f\"Total objects: {total_objects}\\n\"\n stats_text += f\"Average objects per image: {avg_objects_per_image:.2f}\\n\\n\"\n stats_text += \"Class distribution:\\n\"\n for class_name, count in class_distribution.items():\n stats_text += f\" {class_name}: {count}\\n\"\n \n self.text_browser.setPlainText(stats_text)\n \n except Exception as e:\n self.text_browser.setPlainText(f\"An error occurred while generating statistics: {str(e)}\")\n self.show_plot_button.setEnabled(False)\n\n def show_interactive_plot(self):\n if self.plot_file and os.path.exists(self.plot_file):\n webbrowser.open('file://' + os.path.realpath(self.plot_file))\n else:\n self.text_browser.append(\"Error: Plot file not found.\")\n\n def closeEvent(self, event):\n if self.plot_file and os.path.exists(self.plot_file):\n os.unlink(self.plot_file)\n super().closeEvent(event)\n\ndef show_annotation_statistics(parent, annotations):\n dialog = AnnotationStatisticsDialog(parent)\n dialog.generate_statistics(annotations)\n dialog.show_centered(parent)\n return dialog"
},
{
"path": "src/digitalsreeni_image_annotator/annotation_utils.py",
"content": "from PyQt5.QtWidgets import QListWidgetItem\nfrom PyQt5.QtGui import QColor\nfrom PyQt5.QtCore import Qt\n\nclass AnnotationUtils:\n @staticmethod\n def update_annotation_list(self, image_name=None):\n self.annotation_list.clear()\n current_name = image_name or self.current_slice or self.image_file_name\n annotations = self.all_annotations.get(current_name, {})\n for class_name, class_annotations in annotations.items():\n color = self.image_label.class_colors.get(class_name, QColor(Qt.white))\n for i, annotation in enumerate(class_annotations, start=1):\n item_text = f\"{class_name} - {i}\"\n item = QListWidgetItem(item_text)\n item.setData(Qt.UserRole, annotation)\n item.setForeground(color)\n self.annotation_list.addItem(item)\n\n @staticmethod\n def update_slice_list_colors(self):\n for i in range(self.slice_list.count()):\n item = self.slice_list.item(i)\n slice_name = item.text()\n if slice_name in self.all_annotations and any(self.all_annotations[slice_name].values()):\n item.setForeground(QColor(Qt.green))\n else:\n item.setForeground(QColor(Qt.black) if not self.dark_mode else QColor(Qt.white))\n\n @staticmethod\n def update_annotation_list_colors(self, class_name=None, color=None):\n for i in range(self.annotation_list.count()):\n item = self.annotation_list.item(i)\n annotation = item.data(Qt.UserRole)\n if class_name is None or annotation['category_name'] == class_name:\n item_color = color if class_name else self.image_label.class_colors.get(annotation['category_name'], QColor(Qt.white))\n item.setForeground(item_color)\n\n @staticmethod\n def load_image_annotations(self):\n self.image_label.annotations.clear()\n current_name = self.current_slice or self.image_file_name\n if current_name in self.all_annotations:\n self.image_label.annotations = self.all_annotations[current_name].copy()\n self.image_label.update()\n\n @staticmethod\n def save_current_annotations(self):\n current_name = self.current_slice or self.image_file_name\n if current_name:\n if self.image_label.annotations:\n self.all_annotations[current_name] = self.image_label.annotations.copy()\n elif current_name in self.all_annotations:\n del self.all_annotations[current_name]\n AnnotationUtils.update_slice_list_colors(self)\n\n @staticmethod\n def add_annotation_to_list(self, annotation):\n class_name = annotation['category_name']\n color = self.image_label.class_colors.get(class_name, QColor(Qt.white))\n annotations = self.image_label.annotations.get(class_name, [])\n item_text = f\"{class_name} - {len(annotations)}\"\n item = QListWidgetItem(item_text)\n item.setData(Qt.UserRole, annotation)\n item.setForeground(color)\n self.annotation_list.addItem(item)"
},
{
"path": "src/digitalsreeni_image_annotator/annotator_window.py",
"content": "import os\nimport json\nfrom PyQt5.QtWidgets import (QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, \n QPushButton, QFileDialog, QListWidget, QInputDialog, \n QLabel, QButtonGroup, QListWidgetItem, QScrollArea, QCheckBox,\n QSlider, QMenu, QMessageBox, QColorDialog, QDialog, QDoubleSpinBox,\n QGridLayout, QComboBox, QAbstractItemView, QProgressDialog,\n QApplication, QAction, QLineEdit, QTextEdit, QDialogButtonBox, QProgressBar)\nfrom PyQt5.QtGui import QPixmap, QColor, QIcon, QImage, QFont, QKeySequence, QPalette\nfrom PyQt5.QtCore import Qt, QThread, pyqtSignal\nimport numpy as np\nfrom tifffile import TiffFile\nfrom czifile import CziFile\nimport cv2\nfrom datetime import datetime\n\n\nfrom .image_label import ImageLabel\nfrom .utils import calculate_area, calculate_bbox\nfrom .help_window import HelpWindow\n\nfrom .soft_dark_stylesheet import soft_dark_stylesheet\nfrom .default_stylesheet import default_stylesheet\n\nfrom .dataset_splitter import DatasetSplitterTool\nfrom .annotation_statistics import show_annotation_statistics\nfrom .coco_json_combiner import show_coco_json_combiner\nfrom .stack_to_slices import show_stack_to_slices\nfrom .image_patcher import show_image_patcher\nfrom .image_augmenter import show_image_augmenter\nfrom .slice_registration import SliceRegistrationTool\nfrom .sam_utils import SAMUtils\nfrom .snake_game import SnakeGame\nfrom .yolo_trainer import YOLOTrainer, TrainingInfoDialog, LoadPredictionModelDialog\nfrom .stack_interpolator import StackInterpolator\nfrom .dicom_converter import DicomConverter\n\nfrom shapely.geometry import Polygon, MultiPolygon, Point\nfrom shapely.ops import unary_union\nfrom shapely.validation import make_valid\nimport shapely\n\nfrom .export_formats import (\n export_coco_json, export_yolo_v4, export_yolo_v5plus, export_labeled_images,\n export_semantic_labels, export_pascal_voc_bbox, export_pascal_voc_both\n)\n\nfrom .import_formats import import_coco_json, import_yolo_v4, import_yolo_v5plus\nfrom .import_formats import process_import_format\n\nimport shutil \nimport copy\nfrom ultralytics import SAM\n\nimport warnings\nwarnings.filterwarnings(\"ignore\", category=UserWarning)\n\n\nclass TrainingThread(QThread):\n progress_update = pyqtSignal(str)\n finished = pyqtSignal(object)\n\n def __init__(self, yolo_trainer, epochs, imgsz):\n super().__init__()\n self.yolo_trainer = yolo_trainer\n self.epochs = epochs\n self.imgsz = imgsz\n\n def run(self):\n try:\n results = self.yolo_trainer.train_model(epochs=self.epochs, imgsz=self.imgsz)\n self.finished.emit(results)\n except Exception as e:\n self.finished.emit(str(e))\n\nclass DimensionDialog(QDialog):\n def __init__(self, shape, file_name, parent=None, default_dimensions=None):\n super().__init__(parent)\n self.setWindowTitle(\"Assign Dimensions\")\n layout = QVBoxLayout(self)\n \n # Add file name label\n file_name_label = QLabel(f\"File: {file_name}\")\n file_name_label.setWordWrap(True)\n layout.addWidget(file_name_label)\n \n # Add dimension assignment widgets\n dim_widget = QWidget()\n dim_layout = QGridLayout(dim_widget)\n self.combos = []\n self.shape = shape\n dimensions = ['T', 'Z', 'C', 'S', 'H', 'W']\n for i, dim in enumerate(shape):\n dim_layout.addWidget(QLabel(f\"Dimension {i} (size {dim}):\"), i, 0)\n combo = QComboBox()\n combo.addItems(dimensions)\n if default_dimensions and i < len(default_dimensions):\n combo.setCurrentText(default_dimensions[i])\n dim_layout.addWidget(combo, i, 1)\n self.combos.append(combo)\n layout.addWidget(dim_widget)\n \n self.button = QPushButton(\"OK\")\n self.button.clicked.connect(self.accept)\n layout.addWidget(self.button)\n \n self.setMinimumWidth(300)\n\n def get_dimensions(self):\n return [combo.currentText() for combo in self.combos]\n \n\n\nclass ImageAnnotator(QMainWindow):\n def __init__(self):\n super().__init__()\n \n self.is_loading_project = False\n self.backup_project_path = None\n \n self.setWindowTitle(\"Image Annotator\")\n self.setGeometry(100, 100, 1400, 800)\n \n self.central_widget = QWidget()\n self.setCentralWidget(self.central_widget)\n self.layout = QHBoxLayout(self.central_widget)\n \n self.create_menu_bar()\n \n # Initialize image_label early\n self.image_label = ImageLabel()\n self.image_label.set_main_window(self)\n \n # Initialize attributes\n self.current_image = None\n self.current_class = None\n self.image_file_name = \"\"\n self.all_annotations = {}\n self.all_images = []\n self.image_paths = {}\n self.loaded_json = None\n self.class_mapping = {}\n self.editing_mode = False\n self.current_slice = None\n self.slices = []\n self.current_stack = None\n self.image_dimensions = {}\n self.image_slices = {}\n self.image_shapes = {}\n \n # For paint brush and eraser\n self.paint_brush_size = 10\n self.eraser_size = 10\n # Initialize SAM utils\n self.current_sam_model = None\n self.sam_utils = SAMUtils()\n \n # Create sam_magic_wand_button\n self.sam_magic_wand_button = QPushButton(\"Magic Wand\")\n self.sam_magic_wand_button.setCheckable(True)\n self.sam_magic_wand_button.setEnabled(False) # Initially disable the button\n \n # Initialize tool group\n self.tool_group = QButtonGroup(self)\n self.tool_group.setExclusive(False)\n \n # Font size control\n self.font_sizes = {\"Small\": 8, \"Medium\": 10, \"Large\": 12, \"XL\": 14, \"XXL\": 16} # Also, add the options in create_menu_bar method\n self.current_font_size = \"Medium\"\n \n # Dark mode control\n self.dark_mode = False\n \n # Default annotations sorting\n self.current_sort_method = \"class\" # Default sorting method\n \n # Setup UI components\n self.setup_ui()\n \n # Apply theme and font (this includes stylesheet and font size application)\n self.apply_theme_and_font()\n \n # Connect sam_magic_wand_button\n self.sam_magic_wand_button.clicked.connect(self.toggle_tool)\n \n self.class_list.itemChanged.connect(self.toggle_class_visibility)\n \n # YOLO Trainer\n self.yolo_trainer = None\n self.setup_yolo_menu()\n \n # Start in maximized mode\n self.showMaximized()\n\n\n def setup_ui(self):\n # Initialize the main layout\n self.central_widget = QWidget()\n self.setCentralWidget(self.central_widget)\n self.layout = QHBoxLayout(self.central_widget)\n \n # Initialize tool group\n self.tool_group = QButtonGroup(self)\n self.tool_group.setExclusive(False)\n \n # Setup UI components\n self.setup_sidebar()\n self.setup_image_area()\n self.setup_image_list()\n self.setup_slice_list()\n self.update_ui_for_current_tool() \n\n\n\n def update_window_title(self):\n base_title = \"Image Annotator\"\n if hasattr(self, 'current_project_file'):\n project_name = os.path.basename(self.current_project_file)\n project_name = os.path.splitext(project_name)[0] # Remove the file extension\n self.setWindowTitle(f\"{base_title} - {project_name}\")\n else:\n self.setWindowTitle(base_title)\n \n\n \n def new_project(self):\n self.remove_all_temp_annotations() # Remove temp annotations from the previous project\n project_file, _ = QFileDialog.getSaveFileName(self, \"Create New Project\", \"\", \"Image Annotator Project (*.iap)\")\n if project_file:\n # Ensure the file has the correct extension\n if not project_file.lower().endswith('.iap'):\n project_file += '.iap'\n \n self.current_project_file = project_file\n self.current_project_dir = os.path.dirname(project_file)\n \n # Create the images directory\n images_dir = os.path.join(self.current_project_dir, \"images\")\n os.makedirs(images_dir, exist_ok=True)\n \n # Clear existing data without showing messages\n self.clear_all(new_project=True, show_messages=False)\n \n # Prompt for initial project notes\n notes, ok = QInputDialog.getMultiLineText(self, \"Project Notes\", \"Enter initial project notes:\")\n if ok:\n self.project_notes = notes\n else:\n self.project_notes = \"\"\n \n self.project_creation_date = datetime.now().isoformat()\n \n # Save the empty project without showing a message\n self.save_project(show_message=False)\n \n # Keep only this message\n self.show_info(\"New Project\", f\"New project created at {self.current_project_file}\")\n self.initialize_yolo_trainer()\n self.update_window_title()\n \n def show_project_search(self):\n from .project_search import show_project_search\n show_project_search(self) \n \n\n \n def open_project(self):\n print(\"open_project method called\") # Debug print\n self.remove_all_temp_annotations() # Remove temp annotations from the previous project\n project_file, _ = QFileDialog.getOpenFileName(self, \"Open Project\", \"\", \"Image Annotator Project (*.iap)\")\n print(f\"Selected project file: {project_file}\") # Debug print\n if project_file:\n try:\n self.backup_project_before_open(project_file)\n self.open_specific_project(project_file)\n except Exception as e:\n self.restore_project_from_backup()\n QMessageBox.critical(self, \"Error\", f\"An error occurred while opening the project: {str(e)}\\n\"\n f\"The project file has been restored from backup.\")\n else:\n print(\"No project file selected\") # Debug print\n\n\n def backup_project_before_open(self, project_file):\n \"\"\"Create a backup of the project file before opening it.\"\"\"\n import shutil\n import os\n timestamp = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n backup_dir = os.path.join(os.path.dirname(project_file), \".project_backups\")\n os.makedirs(backup_dir, exist_ok=True)\n \n self.backup_project_path = os.path.join(backup_dir, \n f\"{os.path.basename(project_file)}.{timestamp}.backup\")\n shutil.copy2(project_file, self.backup_project_path)\n\n def restore_project_from_backup(self):\n \"\"\"Restore the project file from its backup if available.\"\"\"\n if self.backup_project_path and os.path.exists(self.backup_project_path):\n try:\n shutil.copy2(self.backup_project_path, self.current_project_file)\n print(f\"Project restored from backup: {self.backup_project_path}\")\n except Exception as e:\n print(f\"Failed to restore from backup: {str(e)}\")\n \n \n\n def open_specific_project(self, project_file):\n print(f\"Opening specific project: {project_file}\") # Debug print\n if os.path.exists(project_file):\n try:\n self.is_loading_project = True # Set loading flag\n \n with open(project_file, 'r') as f:\n project_data = json.load(f)\n \n self.clear_all(show_messages=False)\n self.current_project_file = project_file\n self.current_project_dir = os.path.dirname(project_file)\n \n # Load project notes and metadata\n self.project_notes = project_data.get('notes', '')\n self.project_creation_date = project_data.get('creation_date', '')\n self.last_modified = project_data.get('last_modified', '')\n \n # Parse dates\n if self.project_creation_date:\n self.project_creation_date = datetime.fromisoformat(self.project_creation_date).strftime(\"%Y-%m-%d %H:%M:%S\")\n if self.last_modified:\n self.last_modified = datetime.fromisoformat(self.last_modified).strftime(\"%Y-%m-%d %H:%M:%S\")\n \n # Load all data without triggering auto-saves\n self.load_project_data(project_data)\n \n # Now save once after everything is loaded\n self.is_loading_project = False # Clear loading flag\n self.save_project(show_message=False) # Save once after loading\n \n self.initialize_yolo_trainer() \n self.update_window_title()\n \n print(f\"Project opened successfully: {project_file}\")\n QMessageBox.information(self, \"Project Opened\", f\"Project opened successfully: {os.path.basename(project_file)}\")\n \n except Exception as e:\n self.is_loading_project = False # Make sure to clear flag on error\n raise e\n else:\n print(f\"Project file not found: {project_file}\")\n QMessageBox.critical(self, \"Error\", f\"Project file not found: {project_file}\")\n\n def load_project_data(self, project_data):\n \"\"\"Load project data without triggering auto-saves.\"\"\"\n # Load classes\n self.class_mapping.clear()\n self.image_label.class_colors.clear()\n for class_info in project_data.get('classes', []):\n self.add_class(class_info['name'], QColor(class_info['color']))\n \n # Load images\n self.all_images = project_data.get('images', [])\n self.image_paths = project_data.get('image_paths', {})\n \n # Load all annotations first\n self.all_annotations.clear()\n for image_info in project_data['images']:\n if image_info.get('is_multi_slice', False):\n for slice_info in image_info.get('slices', []):\n self.all_annotations[slice_info['name']] = slice_info['annotations']\n else:\n self.all_annotations[image_info['file_name']] = image_info.get('annotations', {})\n \n # Handle missing images\n missing_images = []\n for image_info in project_data['images']:\n image_path = os.path.join(self.current_project_dir, \"images\", image_info['file_name'])\n \n if not os.path.exists(image_path):\n missing_images.append(image_info['file_name'])\n continue\n \n # Update image_paths\n self.image_paths[image_info['file_name']] = image_path\n \n if image_info.get('is_multi_slice', False):\n dimensions = image_info.get('dimensions', [])\n shape = image_info.get('shape', [])\n self.load_multi_slice_image(image_path, dimensions, shape)\n else:\n self.add_images_to_list([image_path])\n \n\n # Update UI\n self.update_ui()\n \n # Handle missing images if any\n if missing_images:\n self.handle_missing_images(missing_images)\n \n # Select the first image if available\n if self.image_list.count() > 0:\n self.image_list.setCurrentRow(0)\n first_item = self.image_list.item(0)\n if first_item:\n self.switch_image(first_item)\n \n # Select the first class if available\n if self.class_list.count() > 0:\n self.class_list.setCurrentRow(0)\n self.on_class_selected()\n \n\n\n def handle_missing_images(self, missing_images):\n message = \"The following images have annotations but were not found in the project directory:\\n\\n\"\n message += \"\\n\".join(missing_images[:10]) # Show first 10 missing images\n if len(missing_images) > 10:\n message += f\"\\n... and {len(missing_images) - 10} more.\"\n message += \"\\n\\nWould you like to locate these images now?\"\n \n reply = QMessageBox.question(self, \"Missing Images\", message, \n QMessageBox.Yes | QMessageBox.No, QMessageBox.Yes)\n \n if reply == QMessageBox.Yes:\n self.load_missing_images(missing_images)\n else:\n self.remove_missing_images(missing_images)\n \n\n def remove_missing_images(self, missing_images):\n for image_name in missing_images:\n # Remove from all_images\n self.all_images = [img for img in self.all_images if img['file_name'] != image_name]\n \n # Remove from image_paths\n self.image_paths.pop(image_name, None)\n \n # Remove from all_annotations\n self.all_annotations.pop(image_name, None)\n \n # If it's a multi-slice image, remove all related slices\n base_name = os.path.splitext(image_name)[0]\n if base_name in self.image_slices:\n for slice_name, _ in self.image_slices[base_name]:\n self.all_annotations.pop(slice_name, None)\n del self.image_slices[base_name]\n \n self.update_ui()\n QMessageBox.information(self, \"Images Removed\", \n f\"{len(missing_images)} missing images and their annotations have been removed from the project.\")\n\n\n \n \n def prompt_load_missing_images(self, missing_images):\n message = \"The following images have annotations but were not found in the project directory:\\n\\n\"\n message += \"\\n\".join(missing_images[:10]) # Show first 10 missing images\n if len(missing_images) > 10:\n message += f\"\\n... and {len(missing_images) - 10} more.\"\n message += \"\\n\\nWould you like to locate these images now?\"\n \n reply = QMessageBox.question(self, \"Load Missing Images\", message, \n QMessageBox.Yes | QMessageBox.No, QMessageBox.Yes)\n \n if reply == QMessageBox.Yes:\n self.load_missing_images(missing_images)\n \n\n\n def load_missing_images(self, missing_images):\n files, _ = QFileDialog.getOpenFileNames(self, \"Select Missing Images\", \"\", \"Image Files (*.png *.jpg *.bmp *.tif *.tiff *.czi)\")\n if files:\n images_loaded = 0\n for file_path in files:\n file_name = os.path.basename(file_path)\n if file_name in missing_images:\n dst_path = os.path.join(self.current_project_dir, \"images\", file_name)\n shutil.copy2(file_path, dst_path)\n self.image_paths[file_name] = dst_path\n \n # Add the image to all_images if it's not already there\n if not any(img['file_name'] == file_name for img in self.all_images):\n self.all_images.append({\n \"file_name\": file_name,\n \"height\": 0, \n \"width\": 0, \n \"id\": len(self.all_images) + 1,\n \"is_multi_slice\": False\n })\n images_loaded += 1\n missing_images.remove(file_name)\n \n self.update_image_list()\n if images_loaded > 0:\n self.image_list.setCurrentRow(0) # Select the first image\n self.switch_image(self.image_list.item(0)) # Display the first image\n QMessageBox.information(self, \"Images Loaded\", \n f\"Successfully copied and loaded {images_loaded} out of {len(files)} selected images.\")\n \n # If there are still missing images, prompt again\n if missing_images:\n self.prompt_load_missing_images(missing_images)\n\n\n def update_image_list(self):\n self.image_list.clear()\n for image_info in self.all_images:\n self.image_list.addItem(image_info['file_name'])\n\n def select_class(self, index):\n if 0 <= index < self.class_list.count():\n item = self.class_list.item(index)\n self.class_list.setCurrentItem(item)\n self.current_class = item.text()\n print(f\"Selected class: {self.current_class}\")\n else:\n print(\"Invalid class index\")\n \n \n def close_project(self):\n if hasattr(self, 'current_project_file'):\n reply = QMessageBox.question(self, 'Close Project',\n \"Do you want to save the current project before closing?\",\n QMessageBox.Yes | QMessageBox.No | QMessageBox.Cancel)\n \n if reply == QMessageBox.Yes:\n self.remove_all_temp_annotations() # Remove temp annotations before saving\n self.save_project(show_message=False) # Save without showing a message\n elif reply == QMessageBox.Cancel:\n return # User cancelled the operation\n \n # Clear all data\n self.clear_all(new_project=True, show_messages=False)\n \n # Reset project-related attributes\n if hasattr(self, 'current_project_file'):\n del self.current_project_file\n if hasattr(self, 'current_project_dir'):\n del self.current_project_dir\n \n # Update the window title\n self.update_window_title()\n \n \n def delete_selected_class(self):\n selected_items = self.class_list.selectedItems()\n if not selected_items:\n QMessageBox.warning(self, \"No Selection\", \"Please select a class to delete.\")\n return\n \n class_name = selected_items[0].text()\n reply = QMessageBox.question(self, 'Delete Class',\n f\"Are you sure you want to delete the class '{class_name}'?\",\n QMessageBox.Yes | QMessageBox.No, QMessageBox.No)\n if reply == QMessageBox.Yes:\n self.delete_class(class_name) # Sreeni note: Implement this method to handle class deletion\n \n def check_missing_images(self):\n missing_images = [img['file_name'] for img in self.all_images if img['file_name'] not in self.image_paths or not os.path.exists(self.image_paths[img['file_name']])]\n if missing_images:\n self.prompt_load_missing_images(missing_images)\n \n\n def convert_to_serializable(self, obj):\n if isinstance(obj, np.integer):\n return int(obj)\n elif isinstance(obj, np.floating):\n return float(obj)\n elif isinstance(obj, np.ndarray):\n return obj.tolist()\n elif isinstance(obj, list):\n return [self.convert_to_serializable(item) for item in obj]\n elif isinstance(obj, dict):\n return {key: self.convert_to_serializable(value) for key, value in obj.items()}\n else:\n return obj\n \n \n def save_project(self, show_message=True):\n if not hasattr(self, 'current_project_file') or not self.current_project_file:\n self.current_project_file, _ = QFileDialog.getSaveFileName(self, \"Save Project\", \"\", \"Image Annotator Project (*.iap)\")\n if not self.current_project_file:\n return # User cancelled the save dialog\n \n self.current_project_dir = os.path.dirname(self.current_project_file)\n \n # Check if images are in the correct directory structure\n images_dir = os.path.join(self.current_project_dir, \"images\")\n os.makedirs(images_dir, exist_ok=True)\n \n images_to_copy = []\n for file_name, src_path in self.image_paths.items():\n dst_path = os.path.join(images_dir, file_name)\n if os.path.abspath(src_path) != os.path.abspath(dst_path):\n if not os.path.exists(dst_path):\n images_to_copy.append((file_name, src_path, dst_path))\n \n if images_to_copy:\n reply = QMessageBox.question(self, 'Image Directory Structure',\n f\"The project structure requires all images to be in an 'images' subdirectory. \"\n f\"{len(images_to_copy)} images need to be copied to the correct location. \"\n f\"Do you want to copy these images?\",\n QMessageBox.Yes | QMessageBox.No, QMessageBox.Yes)\n \n if reply == QMessageBox.Yes:\n for file_name, src_path, dst_path in images_to_copy:\n try:\n shutil.copy2(src_path, dst_path)\n self.image_paths[file_name] = dst_path\n except Exception as e:\n QMessageBox.warning(self, \"Copy Failed\", f\"Failed to copy {file_name}: {str(e)}\")\n return\n else:\n QMessageBox.warning(self, \"Save Cancelled\", \"Project cannot be saved without the correct directory structure.\")\n return\n\n \n # Prepare image data\n images_data = []\n for image_info in self.all_images:\n file_name = image_info['file_name']\n image_data = {\n 'file_name': file_name,\n 'width': image_info['width'],\n 'height': image_info['height'],\n 'is_multi_slice': image_info['is_multi_slice']\n }\n \n if image_data['is_multi_slice']:\n base_name_without_ext = os.path.splitext(file_name)[0]\n image_data['slices'] = []\n for slice_name, _ in self.image_slices.get(base_name_without_ext, []):\n slice_data = {\n 'name': slice_name,\n 'annotations': self.convert_to_serializable(self.all_annotations.get(slice_name, {}))\n }\n image_data['slices'].append(slice_data)\n \n image_data['dimensions'] = self.convert_to_serializable(self.image_dimensions.get(base_name_without_ext, []))\n image_data['shape'] = self.convert_to_serializable(self.image_shapes.get(base_name_without_ext, []))\n else:\n image_data['annotations'] = {}\n for class_name, annotations in self.all_annotations.get(file_name, {}).items():\n image_data['annotations'][class_name] = [ann.copy() for ann in annotations]\n \n images_data.append(image_data)\n \n # Create project data\n project_data = {\n 'classes': [\n {'name': name, 'color': color.name()} \n for name, color in self.image_label.class_colors.items()\n ],\n 'images': images_data,\n 'image_paths': {k: v for k, v in self.image_paths.items() if os.path.exists(v)},\n 'notes': getattr(self, 'project_notes', ''),\n 'creation_date': getattr(self, 'project_creation_date', datetime.now().isoformat()),\n 'last_modified': datetime.now().isoformat()\n }\n \n # Save project data\n with open(self.current_project_file, 'w') as f:\n json.dump(self.convert_to_serializable(project_data), f, indent=2)\n \n if show_message:\n self.show_info(\"Project Saved\", f\"Project saved to {self.current_project_file}\")\n \n # Update the window title\n self.update_window_title()\n \n # Update image_paths to reflect the correct locations\n for file_name in self.image_paths.keys():\n self.image_paths[file_name] = os.path.join(images_dir, file_name)\n \n\n\n def save_project_as(self):\n new_project_file, _ = QFileDialog.getSaveFileName(self, \"Save Project As\", \"\", \"Image Annotator Project (*.iap)\")\n if new_project_file:\n # Ensure the file has the correct extension\n if not new_project_file.lower().endswith('.iap'):\n new_project_file += '.iap'\n \n # Store the original project file\n original_project_file = getattr(self, 'current_project_file', None)\n \n # Set the new project file as the current one\n self.current_project_file = new_project_file\n self.current_project_dir = os.path.dirname(new_project_file)\n \n # Save the project with the new name\n self.save_project(show_message=False)\n \n # Update the window title\n self.update_window_title()\n \n # Show a success message\n QMessageBox.information(self, \"Project Saved As\", f\"Project saved as:\\n{new_project_file}\")\n \n # If this was originally a new unsaved project, update the original project file\n if original_project_file is None:\n self.current_project_file = new_project_file\n\n def auto_save(self):\n if self.is_loading_project:\n return # Skip auto-save during project loading\n \n if not hasattr(self, 'current_project_file'):\n reply = QMessageBox.question(self, 'No Project', \n \"You need to save the project before auto-saving. Would you like to save now?\",\n QMessageBox.Yes | QMessageBox.No, QMessageBox.Yes)\n if reply == QMessageBox.Yes:\n self.save_project()\n else:\n return\n \n if hasattr(self, 'current_project_file'):\n self.save_project(show_message=False)\n print(\"Project auto-saved.\")\n\n def show_project_details(self):\n if not hasattr(self, 'current_project_file'):\n QMessageBox.warning(self, \"No Project\", \"Please open or create a project first.\")\n return\n \n from .project_details import ProjectDetailsDialog\n from .annotation_statistics import AnnotationStatisticsDialog\n \n # Generate annotation statistics\n stats_dialog = AnnotationStatisticsDialog(self)\n stats_dialog.generate_statistics(self.all_annotations)\n \n dialog = ProjectDetailsDialog(self, stats_dialog)\n \n if dialog.exec_() == QDialog.Accepted:\n if dialog.were_changes_made():\n self.project_notes = dialog.get_notes()\n self.save_project(show_message=False)\n QMessageBox.information(self, \"Project Details\", \"Project details have been updated.\")\n else:\n print(\"No changes made to project details.\")\n \n \n \n def load_multi_slice_image(self, image_path, dimensions=None, shape=None):\n \n file_name = os.path.basename(image_path)\n base_name = os.path.splitext(file_name)[0]\n print(f\"Loading multi-slice image: {image_path}\")\n print(f\"Base name: {base_name}\")\n \n if dimensions and shape:\n print(f\"Using stored dimensions: {dimensions}\")\n print(f\"Using stored shape: {shape}\")\n self.image_dimensions[base_name] = dimensions\n self.image_shapes[base_name] = shape\n if image_path.lower().endswith(('.tif', '.tiff')):\n self.load_tiff(image_path, dimensions, shape)\n elif image_path.lower().endswith('.czi'):\n self.load_czi(image_path, dimensions, shape)\n else:\n print(\"No stored dimensions or shape, loading as new image\")\n if image_path.lower().endswith(('.tif', '.tiff')):\n self.load_tiff(image_path)\n elif image_path.lower().endswith('.czi'):\n self.load_czi(image_path)\n \n print(f\"Loaded multi-slice image: {file_name}\")\n print(f\"Dimensions: {self.image_dimensions.get(base_name, 'Not found')}\")\n print(f\"Shape: {self.image_shapes.get(base_name, 'Not found')}\")\n print(f\"Number of slices: {len(self.slices)}\")\n \n if self.slices:\n self.current_image = self.slices[0][1]\n self.current_slice = self.slices[0][0]\n \n self.update_slice_list()\n self.slice_list.setCurrentRow(0)\n self.activate_slice(self.current_slice)\n print(f\"Activated first slice: {self.current_slice}\")\n else:\n print(\"No slices were loaded\")\n self.current_image = None\n self.current_slice = None\n \n self.update_slice_list()\n self.image_label.update()\n \n # print(f\"Loaded slices: {[slice_name for slice_name, _ in self.slices]}\")\n \n\n \n def activate_sam_magic_wand(self):\n # Uncheck all other tools\n for button in self.tool_group.buttons():\n if button != self.sam_magic_wand_button:\n button.setChecked(False)\n \n # Set the current tool\n self.image_label.current_tool = \"sam_magic_wand\"\n self.image_label.sam_magic_wand_active = True\n self.image_label.setCursor(Qt.CrossCursor)\n \n # Update UI based on the current tool\n self.update_ui_for_current_tool()\n \n # If a class is not selected, select the first one (if available)\n if self.current_class is None and self.class_list.count() > 0:\n self.class_list.setCurrentRow(0)\n self.current_class = self.class_list.currentItem().text()\n elif self.class_list.count() == 0:\n QMessageBox.warning(self, \"No Class Selected\", \"Please add a class before using annotation tools.\")\n self.sam_magic_wand_button.setChecked(False)\n self.deactivate_sam_magic_wand()\n \n def deactivate_sam_magic_wand(self):\n self.image_label.current_tool = None\n self.image_label.sam_magic_wand_active = False\n self.sam_magic_wand_button.setChecked(False)\n self.sam_magic_wand_button.setEnabled(False) # Disable the button\n self.image_label.setCursor(Qt.ArrowCursor)\n \n # Clear any SAM-related temporary data\n self.image_label.sam_bbox = None\n self.image_label.drawing_sam_bbox = False\n self.image_label.temp_sam_prediction = None\n \n # Update UI based on the current tool\n self.update_ui_for_current_tool()\n \n def toggle_sam_assisted(self):\n if not self.current_sam_model:\n QMessageBox.warning(self, \"No SAM Model Selected\", \"Please pick a SAM model before using the SAM-Assisted tool.\")\n self.sam_magic_wand_button.setChecked(False)\n return\n \n if self.sam_magic_wand_button.isChecked():\n self.activate_sam_magic_wand()\n else:\n self.deactivate_sam_magic_wand()\n \n self.image_label.clear_temp_sam_prediction() # Clear temporary prediction\n \n def toggle_sam_magic_wand(self):\n if self.sam_magic_wand_button.isChecked():\n if self.current_class is None:\n QMessageBox.warning(self, \"No Class Selected\", \"Please select a class before using SAM2 Magic Wand.\")\n self.sam_magic_wand_button.setChecked(False)\n return\n self.image_label.setCursor(Qt.CrossCursor)\n self.image_label.sam_magic_wand_active = True\n else:\n self.image_label.setCursor(Qt.ArrowCursor)\n self.image_label.sam_magic_wand_active = False\n self.image_label.sam_bbox = None\n \n self.image_label.clear_temp_sam_prediction() # Clear temporary prediction\n \n\n \n def apply_sam_prediction(self):\n if self.image_label.sam_bbox is None:\n print(\"SAM bbox is None\")\n return\n \n x1, y1, x2, y2 = self.image_label.sam_bbox\n bbox = [min(x1, x2), min(y1, y2), max(x1, x2), max(y1, y2)]\n print(f\"Applying SAM prediction with bbox: {bbox}\")\n \n prediction = self.sam_utils.apply_sam_prediction(self.current_image, bbox)\n \n if prediction:\n temp_annotation = {\n \"segmentation\": prediction[\"segmentation\"],\n \"category_id\": self.class_mapping[self.current_class],\n \"category_name\": self.current_class,\n \"score\": prediction[\"score\"]\n }\n \n self.image_label.temp_sam_prediction = temp_annotation\n self.image_label.update()\n else:\n print(\"Failed to generate prediction\")\n \n # Reset SAM bounding box\n self.image_label.sam_bbox = None\n self.image_label.update()\n \n def accept_sam_prediction(self):\n if self.image_label.temp_sam_prediction:\n new_annotation = self.image_label.temp_sam_prediction\n self.image_label.annotations.setdefault(new_annotation[\"category_name\"], []).append(new_annotation)\n self.add_annotation_to_list(new_annotation)\n self.save_current_annotations()\n self.update_slice_list_colors()\n self.image_label.temp_sam_prediction = None\n self.image_label.update()\n print(\"SAM prediction accepted and added to annotations.\")\n \n def setup_slice_list(self):\n self.slice_list = QListWidget()\n self.slice_list.itemClicked.connect(self.switch_slice)\n self.image_list_layout.addWidget(QLabel(\"Slices:\"))\n self.image_list_layout.addWidget(self.slice_list)\n \n\n\n \n def qimage_to_numpy(self, qimage):\n width = qimage.width()\n height = qimage.height()\n ptr = qimage.bits()\n ptr.setsize(height * width * 4)\n arr = np.frombuffer(ptr, np.uint8).reshape((height, width, 4))\n return arr[:, :, :3] # Slice off the alpha channel\n\n\n def open_images(self):\n file_names, _ = QFileDialog.getOpenFileNames(self, \"Open Images\", \"\", \"Image Files (*.png *.jpg *.bmp *.tif *.tiff *.czi)\")\n if file_names:\n self.image_list.clear()\n self.image_paths.clear()\n self.all_images.clear()\n self.slice_list.clear()\n self.slices.clear()\n self.current_stack = None\n self.current_slice = None\n self.add_images_to_list(file_names)\n \n \n def convert_to_8bit_rgb(self, image_array):\n if image_array.ndim == 2:\n # Grayscale image\n image_8bit = self.normalize_array(image_array)\n return np.stack((image_8bit,) * 3, axis=-1)\n elif image_array.ndim == 3:\n if image_array.shape[2] == 3:\n # Already RGB, just normalize\n return self.normalize_array(image_array)\n elif image_array.shape[2] > 3:\n # Multi-channel image, use first three channels\n rgb_array = image_array[:, :, :3]\n return self.normalize_array(rgb_array)\n \n raise ValueError(f\"Unsupported image shape: {image_array.shape}\")\n \n \n \n def add_images_to_list(self, file_names):\n first_added_item = None\n for file_name in file_names:\n base_name = os.path.basename(file_name)\n if base_name not in self.image_paths:\n image_info = {\n \"file_name\": base_name,\n \"height\": 0,\n \"width\": 0,\n \"id\": len(self.all_images) + 1,\n \"is_multi_slice\": False\n }\n \n # Detect multi-slice images and set dimensions\n if file_name.lower().endswith(('.tif', '.tiff', '.czi')):\n self.load_multi_slice_image(file_name)\n base_name_without_ext = os.path.splitext(base_name)[0]\n if base_name_without_ext in self.image_slices and self.image_slices[base_name_without_ext]:\n first_slice_name, first_slice = self.image_slices[base_name_without_ext][0]\n image_info[\"height\"] = first_slice.height()\n image_info[\"width\"] = first_slice.width()\n image_info[\"is_multi_slice\"] = True\n image_info[\"dimensions\"] = self.image_dimensions.get(base_name_without_ext, [])\n image_info[\"shape\"] = self.image_shapes.get(base_name_without_ext, [])\n else:\n # For regular images\n image = QImage(file_name)\n image_info[\"height\"] = image.height()\n image_info[\"width\"] = image.width()\n \n self.all_images.append(image_info)\n item = QListWidgetItem(base_name)\n self.image_list.addItem(item)\n if first_added_item is None:\n first_added_item = item\n \n # Update image_paths with the original file path\n self.image_paths[base_name] = file_name\n \n if first_added_item:\n self.image_list.setCurrentItem(first_added_item)\n self.switch_image(first_added_item)\n\n if not self.is_loading_project:\n self.auto_save()\n \n\n def update_all_images(self, new_image_info):\n for info in new_image_info:\n if not any(img['file_name'] == info['file_name'] for img in self.all_images):\n self.all_images.append(info)\n\n def closeEvent(self, event):\n if not self.image_label.check_unsaved_changes():\n event.ignore()\n return\n event.accept()\n\n if self.image_label.temp_paint_mask is not None or self.image_label.temp_eraser_mask is not None:\n reply = QMessageBox.question(self, 'Unsaved Changes',\n \"You have unsaved changes. Do you want to save them before closing?\",\n QMessageBox.Yes | QMessageBox.No | QMessageBox.Cancel)\n if reply == QMessageBox.Yes:\n if self.image_label.temp_paint_mask is not None:\n self.image_label.commit_paint_annotation()\n if self.image_label.temp_eraser_mask is not None:\n self.image_label.commit_eraser_changes()\n elif reply == QMessageBox.Cancel:\n event.ignore()\n return\n \n # Perform any other cleanup or saving operations here\n event.accept()\n\n \n def switch_slice(self, item):\n if item is None:\n return\n if not self.image_label.check_unsaved_changes():\n return\n \n # Check for unsaved changes\n if self.image_label.temp_paint_mask is not None or self.image_label.temp_eraser_mask is not None:\n reply = QMessageBox.question(self, 'Unsaved Changes',\n \"You have unsaved changes. Do you want to save them?\",\n QMessageBox.Yes | QMessageBox.No | QMessageBox.Cancel)\n if reply == QMessageBox.Yes:\n if self.image_label.temp_paint_mask is not None:\n self.image_label.commit_paint_annotation()\n if self.image_label.temp_eraser_mask is not None:\n self.image_label.commit_eraser_changes()\n elif reply == QMessageBox.Cancel:\n return\n else:\n self.image_label.discard_paint_annotation()\n self.image_label.discard_eraser_changes()\n \n self.save_current_annotations()\n self.image_label.clear_temp_sam_prediction()\n \n slice_name = item.text()\n for name, qimage in self.slices:\n if name == slice_name:\n self.current_image = qimage\n self.current_slice = name\n self.display_image()\n self.load_image_annotations()\n self.update_annotation_list()\n self.clear_highlighted_annotation()\n self.image_label.highlighted_annotations.clear() # Add this line\n self.image_label.reset_annotation_state()\n self.image_label.clear_current_annotation()\n self.update_image_info()\n break\n \n # Ensure the UI is updated\n self.image_label.update()\n self.update_slice_list_colors()\n \n # Reset zoom level to default (1.0)\n self.set_zoom(1.0)\n\n\n def switch_image(self, item):\n if item is None:\n return\n if not self.image_label.check_unsaved_changes(): \n return\n \n # Store the current item before checking temp annotations\n current_item = self.image_list.currentItem()\n \n if not self.check_temp_annotations():\n # If the user chooses not to discard temp annotations, revert the selection\n self.image_list.setCurrentItem(current_item)\n return\n\n self.save_current_annotations()\n self.image_label.clear_temp_sam_prediction()\n self.image_label.exit_editing_mode()\n\n file_name = item.text()\n print(f\"\\nSwitching to image: {file_name}\")\n\n image_info = next((img for img in self.all_images if img[\"file_name\"] == file_name), None)\n \n if image_info:\n self.image_file_name = file_name\n image_path = self.image_paths.get(file_name)\n \n if not image_path:\n image_path = os.path.join(self.current_project_dir, \"images\", file_name)\n\n if image_path and os.path.exists(image_path):\n if image_info.get('is_multi_slice', False):\n base_name = os.path.splitext(file_name)[0]\n if base_name in self.image_slices:\n self.slices = self.image_slices[base_name]\n if self.slices:\n self.current_image = self.slices[0][1]\n self.current_slice = self.slices[0][0]\n self.update_slice_list()\n self.activate_slice(self.current_slice)\n else:\n self.load_multi_slice_image(image_path, image_info.get('dimensions'), image_info.get('shape'))\n else:\n self.load_regular_image(image_path)\n self.display_image()\n self.clear_slice_list()\n \n self.load_image_annotations()\n self.update_annotation_list()\n self.clear_highlighted_annotation()\n \n self.image_label.highlighted_annotations.clear() \n self.image_label.update()\n self.image_label.reset_annotation_state()\n self.image_label.clear_current_annotation()\n self.update_image_info()\n\n self.adjust_zoom_to_fit()\n else:\n self.current_image = None\n self.image_label.clear()\n self.load_image_annotations()\n self.update_annotation_list()\n self.update_image_info()\n \n self.image_list.setCurrentItem(item)\n self.image_label.update()\n self.update_slice_list_colors()\n else:\n self.current_image = None\n self.current_slice = None\n self.image_label.clear()\n self.update_image_info()\n self.clear_slice_list()\n\n def adjust_zoom_to_fit(self):\n if not self.current_image:\n return\n\n # Get the dimensions of the image and the display area\n image_width = self.current_image.width()\n image_height = self.current_image.height()\n display_width = self.scroll_area.viewport().width()\n display_height = self.scroll_area.viewport().height()\n\n # Calculate and apply the zoom factor to fit the longest side\n zoom_factor = min(display_width / image_width, display_height / image_height)\n self.set_zoom(zoom_factor)\n \n \n def activate_current_slice(self):\n if self.current_slice:\n # Ensure the current slice is selected in the slice list\n items = self.slice_list.findItems(self.current_slice, Qt.MatchExactly)\n if items:\n self.slice_list.setCurrentItem(items[0])\n \n # Load annotations for the current slice\n self.load_image_annotations()\n \n # Update the image label\n self.image_label.update()\n \n # Update the annotation list\n self.update_annotation_list()\n\n def load_image(self, image_path):\n extension = os.path.splitext(image_path)[1].lower()\n if extension in ['.tif', '.tiff']:\n self.load_tiff(image_path)\n elif extension == '.czi':\n self.load_czi(image_path)\n else:\n self.load_regular_image(image_path)\n\n\n\n def load_tiff(self, image_path, dimensions=None, shape=None, force_dimension_dialog=False):\n print(f\"Loading TIFF file: {image_path}\")\n with TiffFile(image_path) as tif:\n print(f\"TIFF tags: {tif.pages[0].tags}\")\n \n # Try to access metadata if available\n try:\n metadata = tif.pages[0].tags['ImageDescription'].value\n print(f\"TIFF metadata: {metadata}\")\n except KeyError:\n print(\"No ImageDescription metadata found\")\n \n # Check if it's a multi-page TIFF\n if len(tif.pages) > 1:\n print(f\"Multi-page TIFF detected. Number of pages: {len(tif.pages)}\")\n # Read all pages into a 3D array\n image_array = tif.asarray()\n else:\n print(\"Single-page TIFF detected.\")\n image_array = tif.pages[0].asarray()\n \n print(f\"Image array shape: {image_array.shape}\")\n print(f\"Image array dtype: {image_array.dtype}\")\n print(f\"Image min: {image_array.min()}, max: {image_array.max()}\")\n \n if dimensions and shape and not force_dimension_dialog:\n # Use stored dimensions and shape\n print(f\"Using stored dimensions: {dimensions}\")\n print(f\"Using stored shape: {shape}\")\n image_array = image_array.reshape(shape)\n else:\n # Process as before for new images or when forcing dimension dialog\n print(\"Processing as new image or forcing dimension dialog.\")\n dimensions = None\n \n self.process_multidimensional_image(image_array, image_path, dimensions, force_dimension_dialog)\n \n def load_czi(self, image_path, dimensions=None, shape=None, force_dimension_dialog=False):\n print(f\"Loading CZI file: {image_path}\")\n with CziFile(image_path) as czi:\n image_array = czi.asarray()\n print(f\"CZI array shape: {image_array.shape}\")\n print(f\"CZI array dtype: {image_array.dtype}\")\n print(f\"CZI array min: {image_array.min()}, max: {image_array.max()}\")\n \n\n \n if dimensions and shape and not force_dimension_dialog:\n # Use stored dimensions and shape\n print(f\"Using stored dimensions: {dimensions}\")\n print(f\"Using stored shape: {shape}\")\n image_array = image_array.reshape(shape)\n else:\n # Process as before for new images or when forcing dimension dialog\n print(\"Processing as new image or forcing dimension dialog.\")\n dimensions = None\n \n self.process_multidimensional_image(image_array, image_path, dimensions, force_dimension_dialog)\n \n \n def load_regular_image(self, image_path):\n self.current_image = QImage(image_path)\n self.slices = []\n self.slice_list.clear()\n self.current_slice = None\n \n def process_multidimensional_image(self, image_array, image_path, dimensions=None, force_dimension_dialog=False):\n file_name = os.path.basename(image_path)\n base_name = os.path.splitext(file_name)[0]\n print(f\"Processing file: {file_name}\")\n print(f\"Image array shape: {image_array.shape}\")\n print(f\"Image array dtype: {image_array.dtype}\")\n \n if dimensions is None or force_dimension_dialog:\n if image_array.ndim > 2:\n default_dimensions = ['Z', 'H', 'W'] if image_array.ndim == 3 else ['T', 'Z', 'H', 'W']\n default_dimensions = default_dimensions[-image_array.ndim:]\n \n # Show a progress dialog\n progress = QProgressDialog(\"Assigning dimensions...\", \"Cancel\", 0, 100, self)\n progress.setWindowModality(Qt.WindowModal)\n progress.setMinimumDuration(0)\n progress.setValue(10)\n QApplication.processEvents()\n \n while True:\n dialog = DimensionDialog(image_array.shape, file_name, self, default_dimensions)\n dialog.setWindowFlags(dialog.windowFlags() & ~Qt.WindowContextHelpButtonHint)\n progress.setValue(50)\n QApplication.processEvents()\n if dialog.exec_():\n dimensions = dialog.get_dimensions()\n print(f\"Assigned dimensions: {dimensions}\")\n if 'H' in dimensions and 'W' in dimensions:\n self.image_dimensions[base_name] = dimensions\n break\n else:\n QMessageBox.warning(self, \"Invalid Dimensions\", \"You must assign both H and W dimensions.\")\n else:\n progress.close()\n return\n progress.setValue(100)\n progress.close()\n else:\n dimensions = ['H', 'W']\n self.image_dimensions[base_name] = dimensions\n \n self.image_shapes[base_name] = image_array.shape\n print(f\"Final assigned dimensions: {self.image_dimensions[base_name]}\")\n print(f\"Image shape: {self.image_shapes[base_name]}\")\n \n if self.image_dimensions[base_name]:\n self.create_slices(image_array, self.image_dimensions[base_name], image_path)\n else:\n rgb_image = self.convert_to_8bit_rgb(image_array)\n self.current_image = self.array_to_qimage(rgb_image)\n self.slices = []\n self.slice_list.clear()\n \n if self.slices:\n self.current_image = self.slices[0][1]\n self.current_slice = self.slices[0][0]\n self.slice_list.setCurrentRow(0)\n self.load_image_annotations()\n self.image_label.update()\n \n self.update_image_info()\n \n # Update UI\n self.update_slice_list()\n self.update_annotation_list()\n self.image_label.update()\n\n \n def create_slices(self, image_array, dimensions, image_path):\n base_name = os.path.splitext(os.path.basename(image_path))[0]\n slices = []\n self.slice_list.clear()\n \n print(f\"Creating slices for {base_name}\")\n print(f\"Dimensions: {dimensions}\")\n print(f\"Image array shape: {image_array.shape}\")\n \n # Create and show progress dialog\n progress = QProgressDialog(\"Loading slices...\", \"Cancel\", 0, 100, self)\n progress.setWindowModality(Qt.WindowModal)\n progress.setMinimumDuration(0) # Show immediately\n \n # Handle 2D images\n if image_array.ndim == 2:\n progress.setValue(50) # Update progress\n QApplication.processEvents() # Allow GUI to update\n normalized_array = self.normalize_array(image_array)\n qimage = self.array_to_qimage(normalized_array)\n slice_name = f\"{base_name}\"\n slices.append((slice_name, qimage))\n self.add_slice_to_list(slice_name)\n else:\n # For 3D or higher dimensional arrays\n slice_indices = [i for i, dim in enumerate(dimensions) if dim not in ['H', 'W']]\n \n total_slices = np.prod([image_array.shape[i] for i in slice_indices])\n for idx, _ in enumerate(np.ndindex(tuple(image_array.shape[i] for i in slice_indices))):\n if progress.wasCanceled():\n break\n \n full_idx = [slice(None)] * len(dimensions)\n for i, val in zip(slice_indices, _):\n full_idx[i] = val\n \n slice_array = image_array[tuple(full_idx)]\n rgb_slice = self.convert_to_8bit_rgb(slice_array)\n qimage = self.array_to_qimage(rgb_slice)\n \n slice_name = f\"{base_name}_{'_'.join([f'{dimensions[i]}{val+1}' for i, val in zip(slice_indices, _)])}\"\n slices.append((slice_name, qimage))\n \n self.add_slice_to_list(slice_name)\n \n # Update progress\n progress_value = int((idx + 1) / total_slices * 100)\n progress.setValue(progress_value)\n QApplication.processEvents() # Allow GUI to update\n \n progress.setValue(100) # Ensure progress reaches 100%\n \n self.image_slices[base_name] = slices\n self.slices = slices\n \n if slices:\n self.current_image = slices[0][1]\n self.current_slice = slices[0][0]\n self.slice_list.setCurrentRow(0)\n \n self.activate_slice(self.current_slice)\n \n slice_info = f\"Total slices: {len(slices)}\"\n for dim, size in zip(dimensions, image_array.shape):\n if dim not in ['H', 'W']:\n slice_info += f\", {dim}: {size}\"\n self.update_image_info(additional_info=slice_info)\n else:\n print(\"No slices were created\")\n \n print(f\"Created {len(slices)} slices for {base_name}\")\n return slices\n\n\n\n def add_slice_to_list(self, slice_name):\n item = QListWidgetItem(slice_name)\n \n if self.dark_mode:\n # Dark mode\n item.setBackground(QColor(40, 40, 40)) # Very dark gray background for all items\n if slice_name in self.all_annotations:\n item.setForeground(QColor(60, 60, 60)) # Dark gray text\n item.setBackground(QColor(173, 216, 230)) # Light blue background\n else:\n item.setForeground(QColor(200, 200, 200)) # Light gray text\n else:\n # Light mode\n item.setBackground(QColor(240, 240, 240)) # Very light gray background for all items\n if slice_name in self.all_annotations:\n item.setForeground(QColor(255, 255, 255)) # White text\n item.setBackground(QColor(70, 130, 180)) # Medium-dark blue background\n else:\n item.setForeground(QColor(0, 0, 0)) # Black text\n \n self.slice_list.addItem(item)\n\n\n\n\n \n def normalize_array(self, array):\n # print(f\"Normalizing array. Shape: {array.shape}, dtype: {array.dtype}\")\n # print(f\"Array min: {array.min()}, max: {array.max()}, mean: {array.mean()}\")\n \n array_float = array.astype(np.float32)\n \n if array.dtype == np.uint16:\n array_normalized = (array_float - array.min()) / (array.max() - array.min())\n elif array.dtype == np.uint8:\n # For 8-bit images, use a simple contrast stretching\n p_low, p_high = np.percentile(array_float, (0, 100)) #Change these to 1, 99 or something to stretch the contrast for visualizing 8 bit images\n array_normalized = np.clip(array_float, p_low, p_high)\n array_normalized = (array_normalized - p_low) / (p_high - p_low)\n else:\n array_normalized = (array_float - array.min()) / (array.max() - array.min())\n \n # Apply gamma correction\n gamma = 1.0 # Adjust this value to fine-tune brightness (> 1 for darker, < 1 for brighter)\n array_normalized = np.power(array_normalized, gamma)\n \n return (array_normalized * 255).astype(np.uint8)\n \n def adjust_contrast(self, image, low_percentile=1, high_percentile=99):\n if image.dtype != np.uint8:\n p_low, p_high = np.percentile(image, (low_percentile, high_percentile))\n image_adjusted = np.clip(image, p_low, p_high)\n image_adjusted = (image_adjusted - p_low) / (p_high - p_low)\n return (image_adjusted * 255).astype(np.uint8)\n return image\n\n\n \n def activate_slice(self, slice_name):\n self.current_slice = slice_name\n self.image_file_name = slice_name\n self.load_image_annotations()\n self.update_annotation_list()\n \n for name, qimage in self.slices:\n if name == slice_name:\n self.current_image = qimage\n self.display_image()\n break\n \n self.image_label.update()\n \n items = self.slice_list.findItems(slice_name, Qt.MatchExactly)\n if items:\n self.slice_list.setCurrentItem(items[0])\n\n \n def array_to_qimage(self, array):\n if array.ndim == 2:\n height, width = array.shape\n return QImage(array.data, width, height, width, QImage.Format_Grayscale8)\n elif array.ndim == 3 and array.shape[2] == 3:\n height, width, _ = array.shape\n bytes_per_line = 3 * width\n return QImage(array.data, width, height, bytes_per_line, QImage.Format_RGB888)\n else:\n raise ValueError(f\"Unsupported array shape {array.shape} for conversion to QImage\")\n\n def update_slice_list(self):\n self.slice_list.clear()\n for slice_name, _ in self.slices:\n item = QListWidgetItem(slice_name)\n if slice_name in self.all_annotations:\n item.setForeground(QColor(Qt.green))\n else:\n item.setForeground(QColor(Qt.black) if not self.dark_mode else QColor(Qt.white))\n self.slice_list.addItem(item)\n \n # Select the current slice\n if self.current_slice:\n items = self.slice_list.findItems(self.current_slice, Qt.MatchExactly)\n if items:\n self.slice_list.setCurrentItem(items[0])\n \n def clear_slice_list(self):\n self.slice_list.clear()\n self.slices = []\n self.current_slice = None\n \n def reset_tool_buttons(self):\n for button in self.tool_group.buttons():\n button.setChecked(False) \n\n def keyPressEvent(self, event):\n # Check if the current focus is on a text editing widget\n focused_widget = QApplication.focusWidget()\n if isinstance(focused_widget, (QLineEdit, QTextEdit)):\n super().keyPressEvent(event)\n return\n \n if event.key() == Qt.Key_F2:\n self.launch_snake_game()\n elif event.key() == Qt.Key_Delete:\n # Handle deletions\n if self.class_list.hasFocus() and self.class_list.currentItem():\n self.delete_class(self.class_list.currentItem())\n elif self.annotation_list.hasFocus() and self.annotation_list.selectedItems():\n self.delete_selected_annotations()\n elif self.image_list.hasFocus() and self.image_list.currentItem():\n self.delete_selected_image()\n elif event.key() == Qt.Key_Up or event.key() == Qt.Key_Down:\n # Handle slice navigation\n if self.slice_list.hasFocus():\n current_row = self.slice_list.currentRow()\n if event.key() == Qt.Key_Up and current_row > 0:\n self.slice_list.setCurrentRow(current_row - 1)\n elif event.key() == Qt.Key_Down and current_row < self.slice_list.count() - 1:\n self.slice_list.setCurrentRow(current_row + 1)\n self.switch_slice(self.slice_list.currentItem())\n else:\n # Pass the event to the parent for default handling\n super().keyPressEvent(event)\n elif event.key() == Qt.Key_Return or event.key() == Qt.Key_Enter:\n # Handle accepting visible temporary classes\n if self.has_visible_temp_classes():\n self.accept_visible_temp_classes()\n else:\n super().keyPressEvent(event)\n elif event.key() == Qt.Key_Escape:\n # Handle rejecting visible temporary classes\n if self.has_visible_temp_classes():\n self.reject_visible_temp_classes()\n else:\n super().keyPressEvent(event)\n else:\n # Pass any other key events to the parent for default handling\n super().keyPressEvent(event)\n\n def has_visible_temp_classes(self):\n for i in range(self.class_list.count()):\n item = self.class_list.item(i)\n if item.text().startswith(\"Temp-\") and item.checkState() == Qt.Checked:\n return True\n return False\n \n def launch_snake_game(self):\n #print(\"Launching Snake game\")\n if not hasattr(self, 'snake_game') or not self.snake_game.isVisible():\n self.snake_game = SnakeGame()\n self.snake_game.show()\n self.snake_game.setFocus()\n \n def import_annotations(self):\n if not self.image_label.check_unsaved_changes(): \n return\n print(\"Starting import_annotations\")\n import_format = self.import_format_selector.currentText()\n print(f\"Import format: {import_format}\")\n \n if import_format == \"COCO JSON\":\n file_name, _ = QFileDialog.getOpenFileName(self, \"Import COCO JSON Annotations\", \"\", \"JSON Files (*.json)\")\n if not file_name:\n print(\"No file selected, returning\")\n return\n \n print(f\"Selected file: {file_name}\")\n json_dir = os.path.dirname(file_name)\n images_dir = os.path.join(json_dir, 'images')\n imported_annotations, image_info = import_coco_json(file_name, self.class_mapping)\n \n elif import_format in [\"YOLO (v4 and earlier)\", \"YOLO (v5+)\"]:\n yaml_file, _ = QFileDialog.getOpenFileName(self, \"Select YOLO Dataset YAML\", \"\", \"YAML Files (*.yaml *.yml)\")\n if not yaml_file:\n print(\"No YAML file selected, returning\")\n return\n \n print(f\"Selected YAML file: {yaml_file}\")\n try:\n imported_annotations, image_info = process_import_format(import_format, yaml_file, self.class_mapping)\n yaml_dir = os.path.dirname(yaml_file)\n if import_format == \"YOLO (v4 and earlier)\":\n images_dir = os.path.join(yaml_dir, 'train', 'images')\n else: # YOLO (v5+)\n images_dir = os.path.join(yaml_dir, 'images', 'train') # Preferring train over val\n except ValueError as e:\n QMessageBox.warning(self, \"Import Error\", str(e))\n return\n \n else:\n QMessageBox.warning(self, \"Unsupported Format\", f\"The selected format '{import_format}' is not implemented for import.\")\n return\n \n print(f\"JSON/YOLO directory: {json_dir if import_format == 'COCO JSON' else os.path.dirname(yaml_file)}\")\n print(f\"Images directory: {images_dir}\")\n print(f\"Imported annotations count: {len(imported_annotations)}\")\n print(f\"Image info count: {len(image_info)}\")\n \n images_loaded = 0\n images_not_found = []\n \n for info in image_info.values():\n print(f\"Processing image: {info['file_name']}\")\n image_path = os.path.join(images_dir, info['file_name'])\n \n if os.path.exists(image_path):\n print(f\"Image found at: {image_path}\")\n self.image_paths[info['file_name']] = image_path\n self.all_images.append({\n \"file_name\": info['file_name'],\n \"height\": info['height'],\n \"width\": info['width'],\n \"id\": info['id'],\n \"is_multi_slice\": False\n })\n images_loaded += 1\n else:\n print(f\"Image not found at: {image_path}\")\n images_not_found.append(info['file_name'])\n \n print(f\"Images loaded: {images_loaded}\")\n print(f\"Images not found: {len(images_not_found)}\")\n \n if images_not_found:\n message = f\"The following {len(images_not_found)} images were not found in the 'images' directory:\\n\\n\"\n message += \"\\n\".join(images_not_found[:10])\n if len(images_not_found) > 10:\n message += f\"\\n... and {len(images_not_found) - 10} more.\"\n message += \"\\n\\nDo you want to proceed and ignore annotations for these missing images?\"\n reply = QMessageBox.question(self, \"Missing Images\", message, \n QMessageBox.Yes | QMessageBox.No, QMessageBox.No)\n \n if reply == QMessageBox.No:\n print(\"Import cancelled due to missing images\")\n QMessageBox.information(self, \"Import Cancelled\", \n \"Import cancelled. Please ensure all images are in the 'images' directory and try again.\")\n return\n \n # Update annotations (only for found images)\n for image_name, annotations in imported_annotations.items():\n if image_name not in self.image_paths:\n continue\n self.all_annotations[image_name] = {}\n for category_name, category_annotations in annotations.items():\n self.all_annotations[image_name][category_name] = []\n for i, ann in enumerate(category_annotations, start=1):\n new_ann = {\n \"segmentation\": ann.get(\"segmentation\"),\n \"bbox\": ann.get(\"bbox\"),\n \"category_id\": ann[\"category_id\"],\n \"category_name\": category_name,\n \"number\": i,\n \"type\": ann.get(\"type\", \"polygon\")\n }\n self.all_annotations[image_name][category_name].append(new_ann)\n \n # Update class mapping and colors\n for annotations in self.all_annotations.values():\n for category_name in annotations.keys():\n if category_name not in self.class_mapping:\n new_id = len(self.class_mapping) + 1\n self.class_mapping[category_name] = new_id\n self.image_label.class_colors[category_name] = QColor(Qt.GlobalColor(new_id % 16 + 7))\n \n print(\"Updating UI\")\n # Update UI\n self.update_class_list()\n self.update_image_list()\n self.update_annotation_list()\n \n # Highlight and display the first image\n if self.image_list.count() > 0:\n self.image_list.setCurrentRow(0)\n self.switch_image(self.image_list.item(0))\n \n # Select the first class if available\n if self.class_list.count() > 0:\n self.class_list.setCurrentRow(0)\n self.on_class_selected()\n \n self.image_label.update()\n \n message = f\"Annotations have been imported successfully from {file_name if import_format == 'COCO JSON' else yaml_file}.\\n\"\n message += f\"{images_loaded} images were loaded from the 'images' directory.\\n\"\n if images_not_found:\n message += f\"Annotations for {len(images_not_found)} missing images were ignored.\"\n \n print(\"Import complete, showing message\")\n QMessageBox.information(self, \"Import Complete\", message)\n self.auto_save() # Auto-save after importing annotations\n \n \n \n def export_annotations(self):\n if not self.image_label.check_unsaved_changes(): \n return\n export_format = self.export_format_selector.currentText()\n \n supported_formats = [\n \"COCO JSON\", \"YOLO (v4 and earlier)\", \"YOLO (v5+)\", \"Labeled Images\", \n \"Semantic Labels\", \"Pascal VOC (BBox)\", \"Pascal VOC (BBox + Segmentation)\"\n ]\n \n if export_format not in supported_formats:\n QMessageBox.warning(self, \"Unsupported Format\", f\"The selected format '{export_format}' is not implemented.\")\n return\n \n if export_format == \"COCO JSON\":\n file_name, _ = QFileDialog.getSaveFileName(self, \"Export COCO JSON Annotations\", \"\", \"JSON Files (*.json)\")\n else:\n file_name = QFileDialog.getExistingDirectory(self, f\"Select Output Directory for {export_format} Export\")\n \n if not file_name:\n return\n \n self.save_current_annotations()\n \n if export_format == \"COCO JSON\":\n output_dir = os.path.dirname(file_name)\n json_filename = os.path.basename(file_name)\n json_file, images_dir = export_coco_json(self.all_annotations, self.class_mapping, \n self.image_paths, self.slices, self.image_slices, \n output_dir, json_filename)\n message = \"Annotations have been exported successfully in COCO JSON format.\\n\"\n message += f\"JSON file: {json_file}\\nImages directory: {images_dir}\"\n \n elif export_format == \"YOLO (v4 and earlier)\":\n labels_dir, yaml_path = export_yolo_v4(self.all_annotations, self.class_mapping, \n self.image_paths, self.slices, self.image_slices, file_name)\n message = \"Annotations have been exported successfully in YOLO (v4 and earlier) format.\\n\"\n message += f\"Labels: {labels_dir}\\nYAML: {yaml_path}\"\n \n elif export_format == \"YOLO (v5+)\":\n output_dir, yaml_path = export_yolo_v5plus(self.all_annotations, self.class_mapping, \n self.image_paths, self.slices, self.image_slices, file_name)\n message = \"Annotations have been exported successfully in YOLO (v5+) format.\\n\"\n message += f\"Output directory: {output_dir}\\nYAML: {yaml_path}\"\n \n elif export_format == \"Labeled Images\":\n labeled_images_dir = export_labeled_images(self.all_annotations, self.class_mapping, \n self.image_paths, self.slices, self.image_slices, file_name)\n message = f\"Labeled images have been exported successfully.\\nLabeled Images: {labeled_images_dir}\\n\"\n message += f\"A class summary has been saved in: {os.path.join(labeled_images_dir, 'class_summary.txt')}\"\n \n elif export_format == \"Semantic Labels\":\n semantic_labels_dir = export_semantic_labels(self.all_annotations, self.class_mapping, \n self.image_paths, self.slices, self.image_slices, file_name)\n message = f\"Semantic labels have been exported successfully.\\nSemantic Labels: {semantic_labels_dir}\\n\"\n message += f\"A class-pixel mapping has been saved in: {os.path.join(semantic_labels_dir, 'class_pixel_mapping.txt')}\"\n \n elif export_format == \"Pascal VOC (BBox)\":\n voc_dir = export_pascal_voc_bbox(self.all_annotations, self.class_mapping, \n self.image_paths, self.slices, self.image_slices, file_name)\n message = \"Annotations have been exported successfully in Pascal VOC format (BBox only).\\n\"\n message += f\"Pascal VOC Annotations: {voc_dir}\"\n \n elif export_format == \"Pascal VOC (BBox + Segmentation)\":\n voc_dir = export_pascal_voc_both(self.all_annotations, self.class_mapping, \n self.image_paths, self.slices, self.image_slices, file_name)\n message = \"Annotations have been exported successfully in Pascal VOC format (BBox + Segmentation).\\n\"\n message += f\"Pascal VOC Annotations: {voc_dir}\"\n \n QMessageBox.information(self, \"Export Complete\", message)\n \n\n def save_slices(self, directory):\n slices_saved = False\n for image_file, image_slices in self.image_slices.items():\n for slice_name, qimage in image_slices:\n if slice_name in self.all_annotations and self.all_annotations[slice_name]:\n file_path = os.path.join(directory, f\"{slice_name}.png\")\n qimage.save(file_path, \"PNG\")\n slices_saved = True\n \n return slices_saved\n\n\n def create_coco_annotation(self, ann, image_id, annotation_id):\n coco_ann = {\n \"id\": annotation_id,\n \"image_id\": image_id,\n \"category_id\": ann[\"category_id\"],\n \"area\": calculate_area(ann),\n \"iscrowd\": 0\n }\n \n if \"segmentation\" in ann:\n coco_ann[\"segmentation\"] = [ann[\"segmentation\"]]\n coco_ann[\"bbox\"] = calculate_bbox(ann[\"segmentation\"])\n elif \"bbox\" in ann:\n coco_ann[\"bbox\"] = ann[\"bbox\"]\n \n return coco_ann\n\n def update_all_annotation_lists(self):\n for image_name in self.all_annotations.keys():\n self.update_annotation_list(image_name)\n self.update_annotation_list() # Update for the current image/slice\n \n def update_annotation_list(self, image_name=None):\n self.annotation_list.clear()\n current_name = image_name or self.current_slice or self.image_file_name\n annotations = self.all_annotations.get(current_name, {})\n for class_name, class_annotations in annotations.items():\n if not class_name.startswith(\"Temp-\"): # Only show non-temporary annotations\n color = self.image_label.class_colors.get(class_name, QColor(Qt.white))\n for annotation in class_annotations:\n number = annotation.get('number', 0)\n area = calculate_area(annotation)\n item_text = f\"{class_name} - {number:<3} Area: {area:.2f}\"\n item = QListWidgetItem(item_text)\n item.setData(Qt.UserRole, annotation)\n item.setForeground(color)\n self.annotation_list.addItem(item)\n \n # Force the annotation list to repaint\n self.annotation_list.repaint()\n \n \n \n def update_slice_list_colors(self):\n # Set the background color of the entire list widget\n if self.dark_mode:\n self.slice_list.setStyleSheet(\"QListWidget { background-color: rgb(40, 40, 40); }\")\n else:\n self.slice_list.setStyleSheet(\"QListWidget { background-color: rgb(240, 240, 240); }\")\n \n for i in range(self.slice_list.count()):\n item = self.slice_list.item(i)\n slice_name = item.text()\n \n if self.dark_mode:\n # Dark mode\n if slice_name in self.all_annotations and any(self.all_annotations[slice_name].values()):\n item.setForeground(QColor(60, 60, 60)) # Dark gray text\n item.setBackground(QColor(173, 216, 230)) # Light blue background\n else:\n item.setForeground(QColor(200, 200, 200)) # Light gray text\n item.setBackground(QColor(40, 40, 40)) # Very dark gray background\n else:\n # Light mode\n if slice_name in self.all_annotations and any(self.all_annotations[slice_name].values()):\n item.setForeground(QColor(255, 255, 255)) # White text\n item.setBackground(QColor(70, 130, 180)) # Medium-dark blue background\n else:\n item.setForeground(QColor(0, 0, 0)) # Black text\n item.setBackground(QColor(240, 240, 240)) # Very light gray background\n \n # Force the list to repaint\n self.slice_list.repaint()\n \n\n def update_annotation_list_colors(self, class_name=None, color=None):\n for i in range(self.annotation_list.count()):\n item = self.annotation_list.item(i)\n annotation = item.data(Qt.UserRole)\n # Update only the item for the specific class if class_name is provided\n if class_name is None or annotation['category_name'] == class_name:\n item_color = color if class_name else self.image_label.class_colors.get(annotation['category_name'], QColor(Qt.white))\n item.setForeground(item_color)\n\n def load_image_annotations(self):\n #print(f\"Loading annotations for: {self.current_slice or self.image_file_name}\")\n self.image_label.annotations.clear()\n current_name = self.current_slice or self.image_file_name\n #print(f\"Current name for annotations: {current_name}\")\n #print(f\"All annotations keys: {list(self.all_annotations.keys())}\")\n if current_name in self.all_annotations:\n self.image_label.annotations = copy.deepcopy(self.all_annotations[current_name])\n #print(f\"Loaded annotations: {self.image_label.annotations}\")\n else:\n print(f\"No annotations found for {current_name}\")\n self.image_label.update()\n\n def save_current_annotations(self):\n if self.current_slice:\n current_name = self.current_slice\n elif self.image_file_name:\n current_name = self.image_file_name\n else:\n #print(\"Error: No current slice or image file name set\")\n return\n \n #print(f\"Saving annotations for: {current_name}\")\n if self.image_label.annotations:\n self.all_annotations[current_name] = self.image_label.annotations.copy()\n #print(f\"Saved {len(self.image_label.annotations)} annotations for {current_name}\")\n elif current_name in self.all_annotations:\n del self.all_annotations[current_name]\n #print(f\"Removed annotations for {current_name}\")\n \n self.update_slice_list_colors()\n \n #print(f\"All annotations now: {self.all_annotations.keys()}\")\n #print(f\"Current slice: {self.current_slice}\")\n #print(f\"Current image_file_name: {self.image_file_name}\")\n \n def setup_class_list(self):\n \"\"\"Set up the class list widget.\"\"\"\n self.class_list = QListWidget()\n self.class_list.setContextMenuPolicy(Qt.CustomContextMenu)\n self.class_list.customContextMenuRequested.connect(self.show_class_context_menu)\n self.class_list.itemClicked.connect(self.on_class_selected)\n self.sidebar_layout.addWidget(QLabel(\"Classes:\"))\n self.sidebar_layout.addWidget(self.class_list)\n\n\n\n def setup_tool_buttons(self):\n \"\"\"Set up the tool buttons with grouped manual and automated tools.\"\"\"\n self.tool_group = QButtonGroup(self)\n self.tool_group.setExclusive(False)\n \n # Create a widget for manual tools\n manual_tools_widget = QWidget()\n manual_layout = QVBoxLayout(manual_tools_widget)\n manual_layout.setSpacing(5)\n \n manual_label = QLabel(\"Manual Tools\")\n manual_label.setAlignment(Qt.AlignCenter)\n manual_layout.addWidget(manual_label)\n \n manual_buttons_layout = QHBoxLayout()\n self.polygon_button = QPushButton(\"Polygon\")\n self.polygon_button.setCheckable(True)\n self.rectangle_button = QPushButton(\"Rectangle\")\n self.rectangle_button.setCheckable(True)\n manual_buttons_layout.addWidget(self.polygon_button)\n manual_buttons_layout.addWidget(self.rectangle_button)\n manual_layout.addLayout(manual_buttons_layout)\n \n self.tool_group.addButton(self.polygon_button)\n self.tool_group.addButton(self.rectangle_button)\n self.polygon_button.clicked.connect(self.toggle_tool)\n self.rectangle_button.clicked.connect(self.toggle_tool)\n \n # Create a widget for automated tools\n automated_tools_widget = QWidget()\n automated_layout = QVBoxLayout(automated_tools_widget)\n automated_layout.setSpacing(5)\n \n automated_label = QLabel(\"Automated Tools\")\n automated_label.setAlignment(Qt.AlignCenter)\n automated_layout.addWidget(automated_label)\n \n automated_buttons_layout = QHBoxLayout()\n self.sam_magic_wand_button = QPushButton(\"Magic Wand\")\n self.sam_magic_wand_button.setCheckable(True)\n automated_buttons_layout.addWidget(self.sam_magic_wand_button)\n automated_layout.addLayout(automated_buttons_layout)\n \n self.tool_group.addButton(self.sam_magic_wand_button)\n self.sam_magic_wand_button.clicked.connect(self.toggle_tool)\n \n # Add the grouped tools to the sidebar layout\n self.sidebar_layout.addWidget(manual_tools_widget)\n self.sidebar_layout.addWidget(automated_tools_widget)\n \n\n \n # Set a fixed size for all buttons to make them smaller\n for button in [self.polygon_button, self.rectangle_button, self.load_sam2_button, self.sam_magic_wand_button]:\n button.setFixedSize(100, 30) \n\n def setup_annotation_list(self):\n \"\"\"Set up the annotation list widget.\"\"\"\n self.annotation_list = QListWidget()\n self.annotation_list.setSelectionMode(QAbstractItemView.ExtendedSelection)\n self.annotation_list.itemSelectionChanged.connect(self.update_highlighted_annotations)\n \n \n \n \n def create_menu_bar(self):\n menu_bar = self.menuBar()\n \n # Project Menu\n project_menu = menu_bar.addMenu(\"&Project\")\n \n new_project_action = QAction(\"&New Project\", self)\n new_project_action.setShortcut(QKeySequence.New)\n new_project_action.triggered.connect(self.new_project)\n project_menu.addAction(new_project_action)\n \n open_project_action = QAction(\"&Open Project\", self)\n open_project_action.setShortcut(QKeySequence.Open)\n open_project_action.triggered.connect(self.open_project)\n project_menu.addAction(open_project_action)\n \n save_project_action = QAction(\"&Save Project\", self)\n save_project_action.setShortcut(QKeySequence.Save)\n save_project_action.triggered.connect(self.save_project)\n project_menu.addAction(save_project_action)\n \n save_project_as_action = QAction(\"Save Project &As...\", self)\n save_project_as_action.setShortcut(QKeySequence(\"Ctrl+Shift+S\"))\n save_project_as_action.triggered.connect(self.save_project_as)\n project_menu.addAction(save_project_as_action)\n \n close_project_action = QAction(\"&Close Project\", self)\n close_project_action.setShortcut(QKeySequence(\"Ctrl+W\"))\n close_project_action.triggered.connect(self.close_project)\n project_menu.addAction(close_project_action)\n \n\n project_details_action = QAction(\"Project &Details\", self)\n project_details_action.setShortcut(QKeySequence(\"Ctrl+I\"))\n project_details_action.triggered.connect(self.show_project_details)\n project_menu.addAction(project_details_action)\n \n search_projects_action = QAction(\"&Search Projects\", self)\n search_projects_action.setShortcut(QKeySequence(\"Ctrl+F\"))\n search_projects_action.triggered.connect(self.show_project_search)\n project_menu.addAction(search_projects_action)\n \n # Settings Menu\n settings_menu = menu_bar.addMenu(\"&Settings\")\n \n font_size_menu = settings_menu.addMenu(\"&Font Size\")\n for size in [\"Small\", \"Medium\", \"Large\", \"XL\", \"XXL\"]:\n action = QAction(size, self)\n action.triggered.connect(lambda checked, s=size: self.change_font_size(s))\n font_size_menu.addAction(action)\n \n toggle_dark_mode_action = QAction(\"Toggle &Dark Mode\", self)\n toggle_dark_mode_action.setShortcut(QKeySequence(\"Ctrl+D\"))\n toggle_dark_mode_action.triggered.connect(self.toggle_dark_mode)\n settings_menu.addAction(toggle_dark_mode_action)\n \n # Tools Menu\n tools_menu = menu_bar.addMenu(\"&Tools\")\n \n annotation_stats_action = QAction(\"Annotation Statistics\", self)\n annotation_stats_action.triggered.connect(self.show_annotation_statistics)\n annotation_stats_action.setShortcut(QKeySequence(\"Ctrl+Alt+S\"))\n tools_menu.addAction(annotation_stats_action) \n \n coco_json_combiner_action = QAction(\"COCO JSON Combiner\", self)\n coco_json_combiner_action.triggered.connect(self.show_coco_json_combiner)\n tools_menu.addAction(coco_json_combiner_action)\n \n dataset_splitter_action = QAction(\"Dataset Splitter\", self)\n dataset_splitter_action.triggered.connect(self.open_dataset_splitter)\n tools_menu.addAction(dataset_splitter_action) \n \n stack_to_slices_action = QAction(\"Stack to Slices\", self)\n stack_to_slices_action.triggered.connect(self.show_stack_to_slices)\n tools_menu.addAction(stack_to_slices_action)\n \n image_patcher_action = QAction(\"Image Patcher\", self)\n image_patcher_action.triggered.connect(self.show_image_patcher)\n tools_menu.addAction(image_patcher_action)\n \n image_augmenter_action = QAction(\"Image Augmenter\", self)\n image_augmenter_action.triggered.connect(self.show_image_augmenter)\n tools_menu.addAction(image_augmenter_action)\n \n slice_registration_action = QAction(\"Slice Registration\", self)\n slice_registration_action.triggered.connect(self.show_slice_registration)\n tools_menu.addAction(slice_registration_action)\n\n stack_interpolator_action = QAction(\"Stack Interpolator\", self)\n stack_interpolator_action.triggered.connect(self.show_stack_interpolator)\n tools_menu.addAction(stack_interpolator_action)\n\n dicom_converter_action = QAction(\"DICOM Converter\", self)\n dicom_converter_action.triggered.connect(self.show_dicom_converter)\n tools_menu.addAction(dicom_converter_action)\n \n # Help Menu\n help_menu = menu_bar.addMenu(\"&Help\")\n \n help_action = QAction(\"&Show Help\", self)\n help_action.setShortcut(QKeySequence.HelpContents)\n help_action.triggered.connect(self.show_help)\n help_menu.addAction(help_action)\n \n \n \n\n def change_font_size(self, size):\n self.current_font_size = size\n self.apply_theme_and_font()\n\n def setup_sidebar(self):\n self.sidebar = QWidget()\n self.sidebar_layout = QVBoxLayout(self.sidebar)\n self.layout.addWidget(self.sidebar, 1)\n \n # Helper function to create section headers\n def create_section_header(text):\n label = QLabel(text)\n label.setProperty(\"class\", \"section-header\")\n label.setAlignment(Qt.AlignLeft)\n return label\n \n # Import functionality\n self.import_button = QPushButton(\"Import Annotations with Images\")\n self.import_button.clicked.connect(self.import_annotations)\n self.sidebar_layout.addWidget(self.import_button)\n \n self.import_format_selector = QComboBox()\n self.import_format_selector.addItem(\"COCO JSON\")\n self.import_format_selector.addItem(\"YOLO (v4 and earlier)\") # Modified name\n self.import_format_selector.addItem(\"YOLO (v5+)\") # New format\n \n self.sidebar_layout.addWidget(self.import_format_selector)\n \n # Add spacing\n self.sidebar_layout.addSpacing(20) \n \n self.add_images_button = QPushButton(\"Add New Images\")\n self.add_images_button.clicked.connect(self.add_images)\n self.sidebar_layout.addWidget(self.add_images_button)\n \n self.add_class_button = QPushButton(\"Add Classes\")\n self.add_class_button.clicked.connect(lambda: self.add_class())\n self.sidebar_layout.addWidget(self.add_class_button)\n \n # Class list (without the \"Classes\" header)\n self.class_list = QListWidget()\n self.class_list.setContextMenuPolicy(Qt.CustomContextMenu)\n self.class_list.customContextMenuRequested.connect(self.show_class_context_menu)\n self.class_list.itemClicked.connect(self.on_class_selected)\n self.sidebar_layout.addWidget(self.class_list)\n\n\n button_layout_class_list = QHBoxLayout()\n self.clrButton =QPushButton(self.class_list)\n self.clrButton.setText(\"clear all\")\n self.clrButton.setEnabled(False)\n self.allButton = QPushButton(self.class_list)\n self.allButton.setText(\"select all\")\n self.allButton.setEnabled(False)\n button_layout_class_list.addWidget(self.clrButton)\n button_layout_class_list.addWidget(self.allButton)\n self.clrButton.clicked.connect(lambda : self.toggle_all_class(Qt.Unchecked))\n self.allButton.clicked.connect(lambda : self.toggle_all_class(Qt.Checked))\n self.sidebar_layout.addLayout(button_layout_class_list)\n\n \n # Annotation section\n self.sidebar_layout.addWidget(create_section_header(\"Annotation\"))\n annotation_widget = QWidget()\n annotation_layout = QVBoxLayout(annotation_widget)\n \n # Manual tools subsection\n manual_widget = QWidget()\n manual_layout = QVBoxLayout(manual_widget)\n \n button_layout_top = QHBoxLayout()\n self.polygon_button = QPushButton(\"Polygon\")\n self.polygon_button.setCheckable(True)\n self.rectangle_button = QPushButton(\"Rectangle\")\n self.rectangle_button.setCheckable(True)\n button_layout_top.addWidget(self.polygon_button)\n button_layout_top.addWidget(self.rectangle_button)\n \n button_layout_bottom = QHBoxLayout()\n self.paint_brush_button = QPushButton(\"Paint Brush\")\n self.paint_brush_button.setCheckable(True)\n self.eraser_button = QPushButton(\"Eraser\")\n self.eraser_button.setCheckable(True)\n button_layout_bottom.addWidget(self.paint_brush_button)\n button_layout_bottom.addWidget(self.eraser_button)\n \n manual_layout.addLayout(button_layout_top)\n manual_layout.addLayout(button_layout_bottom)\n \n annotation_layout.addWidget(manual_widget)\n\n # SAM-Assisted tools subsection\n sam_widget = QWidget()\n sam_layout = QVBoxLayout(sam_widget)\n \n # SAM-Assisted button on top\n self.sam_magic_wand_button = QPushButton(\"SAM-Assisted\")\n self.sam_magic_wand_button.setCheckable(True)\n self.sam_magic_wand_button.clicked.connect(self.toggle_sam_assisted)\n sam_layout.addWidget(self.sam_magic_wand_button)\n \n # Add SAM model selector\n self.sam_model_selector = QComboBox()\n self.sam_model_selector.addItem(\"Pick a SAM Model\")\n self.sam_model_selector.addItems(list(self.sam_utils.sam_models.keys()))\n self.sam_model_selector.currentTextChanged.connect(self.change_sam_model)\n sam_layout.addWidget(self.sam_model_selector)\n \n annotation_layout.addWidget(sam_widget)\n \n # Setup tool group\n self.tool_group = QButtonGroup(self)\n self.tool_group.setExclusive(False)\n self.tool_group.addButton(self.polygon_button)\n self.tool_group.addButton(self.rectangle_button)\n self.tool_group.addButton(self.paint_brush_button)\n self.tool_group.addButton(self.eraser_button)\n self.tool_group.addButton(self.sam_magic_wand_button)\n \n self.polygon_button.clicked.connect(self.toggle_tool)\n self.rectangle_button.clicked.connect(self.toggle_tool)\n self.paint_brush_button.clicked.connect(self.toggle_tool)\n self.eraser_button.clicked.connect(self.toggle_tool)\n self.sam_magic_wand_button.clicked.connect(self.toggle_tool)\n \n # Annotations list subsection\n annotation_layout.addWidget(QLabel(\"Annotations\"))\n self.annotation_list = QListWidget()\n self.annotation_list.setSelectionMode(QAbstractItemView.ExtendedSelection)\n self.annotation_list.itemSelectionChanged.connect(self.update_highlighted_annotations)\n annotation_layout.addWidget(self.annotation_list)\n \n # Create a horizontal layout for the sort buttons\n sort_button_layout = QHBoxLayout()\n \n self.sort_by_class_button = QPushButton(\"Sort by Class\")\n self.sort_by_class_button.clicked.connect(self.sort_annotations_by_class)\n sort_button_layout.addWidget(self.sort_by_class_button)\n \n self.sort_by_area_button = QPushButton(\"Sort by Area\")\n self.sort_by_area_button.clicked.connect(self.sort_annotations_by_area)\n sort_button_layout.addWidget(self.sort_by_area_button)\n \n # Add the sort button layout to the annotation layout\n annotation_layout.addLayout(sort_button_layout)\n \n # Delete and Merge annotation buttons\n self.delete_button = QPushButton(\"Delete\")\n self.delete_button.clicked.connect(self.delete_selected_annotations)\n self.merge_button = QPushButton(\"Merge\")\n self.merge_button.clicked.connect(self.merge_annotations)\n self.change_class_button = QPushButton(\"Change Class\")\n self.change_class_button.clicked.connect(self.change_annotation_class)\n \n # Create a horizontal layout for the other buttons\n button_layout = QHBoxLayout()\n button_layout.addWidget(self.delete_button)\n button_layout.addWidget(self.merge_button)\n button_layout.addWidget(self.change_class_button)\n \n # Add the button layout to the annotation layout\n annotation_layout.addLayout(button_layout)\n \n \n # Add export format selector \n self.export_format_selector = QComboBox()\n self.export_format_selector.addItem(\"COCO JSON\")\n self.export_format_selector.addItem(\"YOLO (v4 and earlier)\") # Modified name\n self.export_format_selector.addItem(\"YOLO (v5+)\") # New format\n self.export_format_selector.addItem(\"Labeled Images\")\n self.export_format_selector.addItem(\"Semantic Labels\")\n self.export_format_selector.addItem(\"Pascal VOC (BBox)\")\n self.export_format_selector.addItem(\"Pascal VOC (BBox + Segmentation)\")\n \n annotation_layout.addWidget(QLabel(\"Export Format:\"))\n annotation_layout.addWidget(self.export_format_selector)\n \n self.export_button = QPushButton(\"Export Annotations\")\n self.export_button.clicked.connect(self.export_annotations)\n annotation_layout.addWidget(self.export_button)\n \n # Add the annotation widget to the sidebar\n self.sidebar_layout.addWidget(annotation_widget)\n \n def sort_annotations_by_class(self):\n current_name = self.current_slice or self.image_file_name\n if current_name not in self.all_annotations:\n QMessageBox.information(self, \"No Annotations\", \"There are no annotations to sort for this image.\")\n return\n \n annotations = self.all_annotations[current_name]\n sorted_annotations = []\n for class_name in sorted(annotations.keys()):\n if not class_name.startswith(\"Temp-\"): # Skip temporary classes\n class_annotations = sorted(annotations[class_name], key=lambda x: x.get('number', 0))\n sorted_annotations.extend(class_annotations)\n \n self.update_annotation_list_with_sorted(sorted_annotations)\n \n def sort_annotations_by_area(self):\n current_name = self.current_slice or self.image_file_name\n if current_name not in self.all_annotations:\n QMessageBox.information(self, \"No Annotations\", \"There are no annotations to sort for this image.\")\n return\n \n annotations = self.all_annotations[current_name]\n sorted_annotations = []\n for class_name in annotations.keys():\n if not class_name.startswith(\"Temp-\"): # Skip temporary classes\n class_annotations = sorted(annotations[class_name], key=lambda x: calculate_area(x), reverse=True)\n sorted_annotations.extend(class_annotations)\n \n self.update_annotation_list_with_sorted(sorted_annotations)\n\n def update_annotation_list_with_sorted(self, sorted_annotations):\n self.annotation_list.clear()\n for annotation in sorted_annotations:\n class_name = annotation['category_name']\n if not class_name.startswith(\"Temp-\"): # Only add non-temporary annotations\n number = annotation.get('number', 0)\n area = calculate_area(annotation)\n item_text = f\"{class_name} - {number:<3} Area: {area:.2f}\"\n item = QListWidgetItem(item_text)\n item.setData(Qt.UserRole, annotation)\n color = self.image_label.class_colors.get(class_name, QColor(Qt.white))\n item.setForeground(color)\n self.annotation_list.addItem(item)\n \n self.image_label.update()\n \n \n def change_sam_model(self, model_name):\n self.sam_utils.change_sam_model(model_name)\n self.current_sam_model = self.sam_utils.current_sam_model\n \n if model_name != \"Pick a SAM Model\":\n # Enable the SAM Magic Wand button\n self.sam_magic_wand_button.setEnabled(True)\n \n # Activate the SAM Magic Wand tool\n self.sam_magic_wand_button.setChecked(True)\n self.activate_sam_magic_wand()\n \n print(f\"Changed SAM model to: {model_name}\")\n else:\n # Disable and deactivate the SAM Magic Wand button\n self.sam_magic_wand_button.setEnabled(False)\n self.sam_magic_wand_button.setChecked(False)\n self.deactivate_sam_magic_wand()\n print(\"SAM model unset\")\n \n \n def setup_font_size_selector(self):\n font_size_label = QLabel(\"Font Size:\")\n self.font_size_selector = QComboBox()\n self.font_size_selector.addItems([\"Small\", \"Medium\", \"Large\"])\n self.font_size_selector.setCurrentText(\"Medium\")\n self.font_size_selector.currentTextChanged.connect(self.on_font_size_changed)\n \n self.sidebar_layout.addWidget(font_size_label)\n self.sidebar_layout.addWidget(self.font_size_selector)\n \n def on_font_size_changed(self, size):\n self.current_font_size = size\n self.apply_theme_and_font()\n \n\n \n def apply_theme_and_font(self):\n font_size = self.font_sizes[self.current_font_size]\n if self.dark_mode:\n style = soft_dark_stylesheet\n else:\n style = default_stylesheet \n \n # Combine the theme stylesheet with font size\n combined_style = f\"{style}\\nQWidget {{ font-size: {font_size}pt; }}\"\n self.setStyleSheet(combined_style)\n \n # Apply font size to all widgets\n for widget in self.findChildren(QWidget):\n font = widget.font()\n font.setPointSize(font_size)\n widget.setFont(font)\n \n self.image_label.setFont(QFont(\"Arial\", font_size))\n self.update()\n\n \n def toggle_dark_mode(self):\n self.dark_mode = not self.dark_mode\n self.apply_theme_and_font()\n \n # Update slice list colors\n self.update_slice_list_colors()\n \n # Update other UI elements if necessary\n self.update_class_list()\n self.update_annotation_list()\n \n # Force a repaint of the main window\n self.repaint()\n \n def apply_stylesheet(self):\n if self.dark_mode:\n self.setStyleSheet(soft_dark_stylesheet)\n else:\n self.setStyleSheet(default_stylesheet)\n \n def update_ui_colors(self):\n # Update colors for elements that need to retain their functionality\n self.update_annotation_list_colors()\n self.update_slice_list_colors()\n self.image_label.update()\n \n def setup_image_area(self):\n \"\"\"Set up the main image area.\"\"\"\n self.image_widget = QWidget()\n self.image_layout = QVBoxLayout(self.image_widget)\n self.layout.addWidget(self.image_widget, 3)\n\n self.scroll_area = QScrollArea()\n self.scroll_area.setWidgetResizable(True)\n self.scroll_area.setHorizontalScrollBarPolicy(Qt.ScrollBarAsNeeded)\n self.scroll_area.setVerticalScrollBarPolicy(Qt.ScrollBarAsNeeded)\n \n # Use the already initialized image_label\n self.image_label.setAlignment(Qt.AlignCenter)\n self.scroll_area.setWidget(self.image_label)\n \n self.image_layout.addWidget(self.scroll_area)\n\n\n self.zoom_slider = QSlider(Qt.Horizontal)\n self.zoom_slider.setMinimum(10)\n self.zoom_slider.setMaximum(500)\n self.zoom_slider.setValue(100)\n self.zoom_slider.setTickPosition(QSlider.TicksBelow)\n self.zoom_slider.setTickInterval(50)\n self.zoom_slider.valueChanged.connect(self.zoom_image)\n self.image_layout.addWidget(self.zoom_slider)\n self.image_info_label = QLabel()\n self.image_layout.addWidget(self.image_info_label)\n\n def setup_image_list(self):\n \"\"\"Set up the image list area.\"\"\"\n self.image_list_widget = QWidget()\n self.image_list_layout = QVBoxLayout(self.image_list_widget)\n self.layout.addWidget(self.image_list_widget, 1)\n\n self.image_list_label = QLabel(\"Images:\")\n self.image_list_layout.addWidget(self.image_list_label)\n\n self.image_list = QListWidget()\n self.image_list.itemClicked.connect(self.switch_image)\n self.image_list.currentRowChanged.connect(lambda row: self.switch_image(self.image_list.currentItem()))\n self.image_list.setContextMenuPolicy(Qt.CustomContextMenu)\n self.image_list.customContextMenuRequested.connect(self.show_image_context_menu)\n self.image_list_layout.addWidget(self.image_list)\n\n self.clear_all_button = QPushButton(\"Clear All Images and Annotations\")\n self.clear_all_button.clicked.connect(self.clear_all)\n self.image_list_layout.addWidget(self.clear_all_button)\n\n########## ### Tools ########## I love useful image processing tools :)\n def open_dataset_splitter(self):\n self.dataset_splitter = DatasetSplitterTool(self)\n self.dataset_splitter.setWindowModality(Qt.ApplicationModal)\n self.dataset_splitter.show_centered(self)\n \n def show_annotation_statistics(self):\n if not self.all_annotations:\n QMessageBox.warning(self, \"No Annotations\", \"There are no annotations to analyze.\")\n return\n try:\n self.annotation_stats_dialog = show_annotation_statistics(self, self.all_annotations)\n except Exception as e:\n QMessageBox.critical(self, \"Error\", f\"An error occurred while showing annotation statistics: {str(e)}\")\n\n \n def show_coco_json_combiner(self):\n self.coco_json_combiner_dialog = show_coco_json_combiner(self)\n \n def show_stack_to_slices(self):\n self.stack_to_slices_dialog = show_stack_to_slices(self)\n \n\n def show_image_patcher(self):\n self.image_patcher_dialog = show_image_patcher(self) \n \n def show_image_augmenter(self):\n self.image_augmenter_dialog = show_image_augmenter(self)\n\n def show_slice_registration(self):\n self.slice_registration_dialog = SliceRegistrationTool(self)\n self.slice_registration_dialog.show_centered(self)\n \n def show_stack_interpolator(self):\n self.stack_interpolator_dialog = StackInterpolator(self)\n self.stack_interpolator_dialog.show_centered(self)\n\n def show_dicom_converter(self):\n self.dicom_converter_dialog = DicomConverter(self)\n self.dicom_converter_dialog.show_centered(self)\n\n###################################################################\n\n # update the show_help method:\n def show_help(self):\n self.help_window = HelpWindow(dark_mode=self.dark_mode, font_size=self.font_sizes[self.current_font_size])\n self.help_window.show_centered(self)\n\n \n def add_images(self):\n if not self.image_label.check_unsaved_changes(): \n return\n file_names, _ = QFileDialog.getOpenFileNames(self, \"Add Images\", \"\", \"Image Files (*.png *.jpg *.bmp *.tif *.tiff *.czi)\")\n if file_names:\n self.add_images_to_list(file_names)\n \n \n def clear_all(self, new_project=False, show_messages=True):\n if not new_project and show_messages:\n reply = self.show_question('Clear All',\n \"Are you sure you want to clear all images and annotations? This action cannot be undone.\")\n if reply != QMessageBox.Yes:\n return\n \n # Clear images\n self.image_list.clear()\n self.image_paths.clear()\n self.all_images.clear()\n self.current_image = None\n self.image_file_name = \"\"\n \n # Clear the image display\n self.image_label.clear()\n self.image_label.setPixmap(QPixmap()) # Set an empty pixmap\n self.image_label.original_pixmap = None\n self.image_label.scaled_pixmap = None\n \n # Clear annotations\n self.all_annotations.clear()\n self.annotation_list.clear()\n self.image_label.annotations.clear()\n self.image_label.highlighted_annotations.clear()\n \n # Clear current class\n self.current_class = None\n \n # Reset class-related data\n self.class_list.clear()\n self.allButton.setEnabled(False)\n self.clrButton.setEnabled(False)\n \n self.image_label.class_colors.clear()\n self.class_mapping.clear()\n \n # Clear slices\n self.image_slices.clear()\n self.slices = []\n self.slice_list.clear()\n self.current_slice = None\n self.current_stack = None\n \n # Reset zoom\n self.image_label.zoom_factor = 1.0\n self.zoom_slider.setValue(100)\n \n # Reset tools\n self.image_label.current_tool = None\n self.polygon_button.setChecked(False)\n self.rectangle_button.setChecked(False)\n self.sam_magic_wand_button.setChecked(False)\n self.sam_magic_wand_button.setEnabled(False) # Disable the SAM-Assisted button\n self.image_label.sam_magic_wand_active = False # Deactivate SAM magic wand\n \n # Reset SAM-related attributes\n self.image_label.sam_bbox = None\n self.image_label.drawing_sam_bbox = False\n self.image_label.temp_sam_prediction = None\n \n self.image_label.setCursor(Qt.ArrowCursor) # Reset cursor to default\n self.sam_model_selector.setCurrentIndex(0) # Reset to \"Pick a SAM Model\"\n self.current_sam_model = None # Reset the current SAM model\n \n # Reset project-related attributes\n if not new_project:\n if hasattr(self, 'current_project_file'):\n del self.current_project_file\n if hasattr(self, 'current_project_dir'):\n del self.current_project_dir\n \n # Update UI\n self.image_label.update()\n self.update_image_info()\n\n \n \n # Force a repaint of the main window\n self.repaint()\n self.update_window_title()\n \n \n def show_warning(self, title, message):\n QMessageBox.warning(self, title, message)\n\n def show_info(self, title, message):\n QMessageBox.information(self, title, message)\n \n def update_image_info(self, additional_info=None):\n if self.current_image:\n width = self.current_image.width()\n height = self.current_image.height()\n info = f\"Image: {width}x{height}\"\n if additional_info:\n info += f\", {additional_info}\"\n self.image_info_label.setText(info)\n else:\n self.image_info_label.setText(\"No image loaded\")\n \n \n def show_question(self, title, message):\n return QMessageBox.question(self, title, message,\n QMessageBox.Yes | QMessageBox.No,\n QMessageBox.No)\n \n def show_image_context_menu(self, position):\n menu = QMenu()\n current_item = self.image_list.itemAt(position)\n if current_item:\n file_name = current_item.text()\n delete_action = menu.addAction(\"Remove Image\")\n \n if not self.is_multi_dimensional(file_name):\n predict_action = menu.addAction(\"Predict using YOLO\")\n \n if self.is_multi_dimensional(file_name):\n redefine_dimensions_action = menu.addAction(\"Redefine Dimensions\")\n \n action = menu.exec_(self.image_list.mapToGlobal(position))\n \n if action == delete_action:\n self.remove_image()\n elif not self.is_multi_dimensional(file_name) and action == predict_action:\n self.predict_single_image(file_name)\n elif self.is_multi_dimensional(file_name) and action == redefine_dimensions_action:\n self.redefine_dimensions(file_name)\n \n def is_multi_dimensional(self, file_name):\n return file_name.lower().endswith(('.tif', '.tiff', '.czi'))\n \n def predict_single_image(self, file_name):\n if self.is_multi_dimensional(file_name):\n return # Do nothing for multi-dimensional images\n \n if not self.yolo_trainer or not self.yolo_trainer.model:\n QMessageBox.warning(self, \"No Model\", \"Please load a YOLO model first from the YOLO > Prediction Settings > Load Model menu.\")\n return\n \n # Deactivate SAM tool before prediction\n self.deactivate_sam_magic_wand()\n \n image_path = self.image_paths[file_name]\n try:\n results = self.yolo_trainer.predict(image_path)\n self.process_yolo_results(results, file_name)\n except Exception as e:\n QMessageBox.warning(self, \"Prediction Error\", \n f\"An error occurred during prediction: {str(e)}\\n\\n\"\n \"This might be due to a mismatch between the model and the YAML file classes. \"\n \"Please check that the YAML file corresponds to the loaded model.\")\n \n def redefine_dimensions(self, file_name):\n file_path = self.image_paths.get(file_name)\n if not file_path or not file_path.lower().endswith(('.tif', '.tiff', '.czi')):\n return # Exit the method if it's not a TIFF or CZI file\n \n reply = QMessageBox.warning(self, \"Redefine Dimensions\",\n \"Redefining dimensions will cause all associated annotations to be lost. \"\n \"Do you want to continue?\",\n QMessageBox.Yes | QMessageBox.No, QMessageBox.No)\n \n if reply == QMessageBox.Yes:\n # Remove existing annotations for this file\n base_name = os.path.splitext(file_name)[0]\n \n print(f\"Removing annotations for image: {base_name}\")\n #print(f\"Current annotations: {list(self.all_annotations.keys())}\")\n \n # Create a list of keys to remove, using a more specific matching condition\n keys_to_remove = [key for key in self.all_annotations.keys() \n if key == base_name or (key.startswith(f\"{base_name}_\") and not key.startswith(f\"{base_name}_8bit\"))]\n \n print(f\"Keys to remove: {keys_to_remove}\")\n \n # Remove the annotations\n for key in keys_to_remove:\n del self.all_annotations[key]\n \n #print(f\"Annotations after removal: {list(self.all_annotations.keys())}\")\n \n # Remove existing slices\n if base_name in self.image_slices:\n del self.image_slices[base_name]\n \n # Clear current image if it's the one being redefined\n if self.image_file_name == file_name:\n self.current_image = None\n self.image_label.clear()\n \n # Reload the image with new dimension dialog\n if file_path.lower().endswith(('.tif', '.tiff')):\n self.load_tiff(file_path, force_dimension_dialog=True)\n elif file_path.lower().endswith('.czi'):\n self.load_czi(file_path, force_dimension_dialog=True)\n \n # Update UI\n self.update_slice_list()\n self.update_annotation_list()\n self.image_label.update()\n \n #print(f\"Final annotations: {list(self.all_annotations.keys())}\")\n \n QMessageBox.information(self, \"Dimensions Redefined\", \n \"The dimensions have been redefined and the image reloaded. \"\n \"All previous annotations for this image have been removed.\")\n \n def remove_image(self):\n current_item = self.image_list.currentItem()\n if current_item:\n file_name = current_item.text()\n \n # Remove from all data structures\n self.image_list.takeItem(self.image_list.row(current_item))\n self.image_paths.pop(file_name, None)\n self.all_images = [img for img in self.all_images if img[\"file_name\"] != file_name]\n \n # Remove annotations\n self.all_annotations.pop(file_name, None)\n \n # Handle multi-dimensional images\n base_name = os.path.splitext(file_name)[0]\n if base_name in self.image_slices:\n # Remove slices\n for slice_name, _ in self.image_slices[base_name]:\n self.all_annotations.pop(slice_name, None)\n del self.image_slices[base_name]\n \n # Clear slice list\n self.slice_list.clear()\n \n # Clear current image and slice if it was the removed image\n if self.image_file_name == file_name:\n self.current_image = None\n self.image_file_name = \"\"\n self.current_slice = None\n self.image_label.clear()\n self.annotation_list.clear()\n \n # Switch to another image if available\n if self.image_list.count() > 0:\n next_item = self.image_list.item(0)\n self.image_list.setCurrentItem(next_item)\n self.switch_image(next_item)\n else:\n # No images left\n self.current_image = None\n self.image_file_name = \"\"\n self.current_slice = None\n self.image_label.clear()\n self.annotation_list.clear()\n self.slice_list.clear()\n \n # Update UI\n self.update_ui() \n self.auto_save() # Auto-save after removing an image\n\n\n def load_annotations(self):\n file_name, _ = QFileDialog.getOpenFileName(self, \"Load Annotations\", \"\", \"JSON Files (*.json)\")\n if file_name:\n with open(file_name, 'r') as f:\n self.loaded_json = json.load(f)\n \n # Load categories\n self.class_list.clear()\n self.image_label.class_colors.clear()\n self.class_mapping.clear()\n for category in self.loaded_json[\"categories\"]:\n class_name = category[\"name\"]\n self.class_mapping[class_name] = category[\"id\"]\n \n # Assign a color if not already assigned\n if class_name not in self.image_label.class_colors:\n color = QColor(Qt.GlobalColor(len(self.image_label.class_colors) % 16 + 7))\n self.image_label.class_colors[class_name] = color\n \n # Add item to class list with color indicator\n item = QListWidgetItem(class_name)\n self.update_class_item_color(item, self.image_label.class_colors[class_name])\n self.class_list.addItem(item)\n \n # Create a mapping of image IDs to file names\n image_id_to_filename = {img[\"id\"]: img[\"file_name\"] for img in self.loaded_json[\"images\"]}\n \n # Load image information\n json_images = {img[\"file_name\"]: img for img in self.loaded_json[\"images\"]}\n \n # Update existing images and add new ones from JSON\n updated_all_images = []\n for i in range(self.image_list.count()):\n item = self.image_list.item(i)\n file_name = item.text()\n if file_name in json_images:\n updated_image = self.all_images[i].copy()\n updated_image.update(json_images[file_name])\n updated_all_images.append(updated_image)\n del json_images[file_name]\n else:\n updated_all_images.append(self.all_images[i])\n \n # Add remaining images from JSON\n for img in json_images.values():\n updated_all_images.append(img)\n self.image_list.addItem(img[\"file_name\"])\n \n self.all_images = updated_all_images\n \n # Load annotations\n self.all_annotations.clear()\n for annotation in self.loaded_json[\"annotations\"]:\n image_id = annotation[\"image_id\"]\n file_name = image_id_to_filename.get(image_id)\n if file_name:\n if file_name not in self.all_annotations:\n self.all_annotations[file_name] = {}\n \n category = next((cat for cat in self.loaded_json[\"categories\"] if cat[\"id\"] == annotation[\"category_id\"]), None)\n if category:\n category_name = category[\"name\"]\n if category_name not in self.all_annotations[file_name]:\n self.all_annotations[file_name][category_name] = []\n \n ann = {\n \"category_id\": annotation[\"category_id\"],\n \"category_name\": category_name,\n }\n \n if \"segmentation\" in annotation:\n ann[\"segmentation\"] = annotation[\"segmentation\"][0]\n ann[\"type\"] = \"polygon\"\n elif \"bbox\" in annotation:\n ann[\"bbox\"] = annotation[\"bbox\"]\n ann[\"type\"] = \"bbox\"\n \n # Add number field if it's missing\n if \"number\" not in ann:\n ann[\"number\"] = len(self.all_annotations[file_name][category_name]) + 1\n \n self.all_annotations[file_name][category_name].append(ann)\n \n # Check for missing images\n missing_images = [img[\"file_name\"] for img in self.loaded_json[\"images\"] if img[\"file_name\"] not in self.image_paths]\n if missing_images:\n self.show_warning(\"Missing Images\", \"The following images are missing:\\n\" + \"\\n\".join(missing_images))\n \n # Reload the current image if it exists, otherwise load the first image\n if self.image_file_name and self.image_file_name in self.all_annotations:\n self.switch_image(self.image_list.findItems(self.image_file_name, Qt.MatchExactly)[0])\n elif self.all_images:\n self.switch_image(self.image_list.item(0))\n \n self.image_label.highlighted_annotations = [] # Clear existing highlights\n self.update_annotation_list() # This will repopulate the annotation list\n self.image_label.update() # Force a redraw of the image label\n\n if self.class_list.count() > 0:\n self.allButton.setEnabled(True)\n self.clrButton.setEnabled(True)\n\n\n def clear_highlighted_annotation(self):\n self.image_label.highlighted_annotation = None\n self.image_label.update()\n \n def update_highlighted_annotations(self):\n selected_items = self.annotation_list.selectedItems()\n self.image_label.highlighted_annotations = [item.data(Qt.UserRole) for item in selected_items]\n self.image_label.update() # Force a redraw of the image label\n \n # Enable/disable merge and change class buttons based on selection\n self.merge_button.setEnabled(len(selected_items) >= 2)\n self.change_class_button.setEnabled(len(selected_items) > 0)\n\n def renumber_annotations(self):\n current_name = self.current_slice or self.image_file_name\n if current_name in self.all_annotations:\n for class_name, annotations in self.all_annotations[current_name].items():\n for i, ann in enumerate(annotations, start=1):\n ann['number'] = i\n self.update_annotation_list()\n\n def delete_selected_annotations(self):\n selected_items = self.annotation_list.selectedItems()\n if not selected_items:\n QMessageBox.warning(self, \"No Selection\", \"Please select an annotation to delete.\")\n return\n \n reply = QMessageBox.question(self, 'Delete Annotations',\n f\"Are you sure you want to delete {len(selected_items)} annotation(s)?\",\n QMessageBox.Yes | QMessageBox.No, QMessageBox.No)\n if reply == QMessageBox.Yes:\n # Create a list of annotations to remove\n annotations_to_remove = []\n for item in selected_items:\n annotation = item.data(Qt.UserRole)\n annotations_to_remove.append((annotation['category_name'], annotation))\n \n # Remove annotations from image_label.annotations\n for category_name, annotation in annotations_to_remove:\n if category_name in self.image_label.annotations:\n if annotation in self.image_label.annotations[category_name]:\n self.image_label.annotations[category_name].remove(annotation)\n \n # Update all_annotations\n current_name = self.current_slice or self.image_file_name\n self.all_annotations[current_name] = self.image_label.annotations\n \n # Sort and update the annotation list based on the current sorting method\n if self.current_sort_method == \"area\":\n self.sort_annotations_by_area()\n else:\n self.sort_annotations_by_class()\n \n self.image_label.highlighted_annotations.clear()\n self.image_label.update()\n \n # Update slice list colors\n self.update_slice_list_colors()\n \n QMessageBox.information(self, \"Annotations Deleted\", f\"{len(selected_items)} annotation(s) have been deleted.\") \n self.auto_save() # Auto-save after deleting annotations\n\n\n def merge_annotations(self):\n if self.image_label.editing_polygon is not None:\n QMessageBox.warning(self, \"Edit Mode Active\", \n \"Please exit the annotation edit mode before merging annotations.\")\n return\n \n selected_items = self.annotation_list.selectedItems()\n if len(selected_items) < 2:\n QMessageBox.warning(self, \"Not Enough Annotations\", \"Please select at least two annotations to merge.\")\n return\n \n class_name = selected_items[0].data(Qt.UserRole)['category_name']\n if not all(item.data(Qt.UserRole)['category_name'] == class_name for item in selected_items):\n QMessageBox.warning(self, \"Mixed Classes\", \"All selected annotations must be from the same class.\")\n return\n \n polygons = []\n original_annotations = []\n for item in selected_items:\n annotation = item.data(Qt.UserRole)\n original_annotations.append(annotation)\n if 'segmentation' in annotation:\n points = zip(annotation['segmentation'][0::2], annotation['segmentation'][1::2])\n polygon = Polygon(points)\n if not polygon.is_valid:\n polygon = polygon.buffer(0)\n polygons.append(polygon)\n \n def are_all_polygons_connected(polygons):\n if len(polygons) < 2:\n return True\n \n connected = set([0]) # Start with the first polygon\n to_check = set(range(1, len(polygons)))\n \n while to_check:\n newly_connected = set()\n for i in connected:\n for j in to_check:\n if polygons[i].intersects(polygons[j]) or polygons[i].touches(polygons[j]):\n newly_connected.add(j)\n \n if not newly_connected:\n return False # If no new connections found, they're not all connected\n \n connected.update(newly_connected)\n to_check -= newly_connected\n \n return True # All polygons are connected\n \n if not are_all_polygons_connected(polygons):\n QMessageBox.warning(self, \"Disconnected Polygons\", \"Not all selected annotations are connected. Please select only connected annotations to merge.\")\n return\n \n try:\n merged_polygon = unary_union(polygons)\n except Exception as e:\n QMessageBox.warning(self, \"Merge Error\", f\"Unable to merge the selected annotations due to an error: {str(e)}\")\n return\n \n new_annotation = {\n \"segmentation\": [],\n \"category_id\": self.class_mapping[class_name],\n \"category_name\": class_name,\n }\n \n if isinstance(merged_polygon, Polygon):\n new_annotation[\"segmentation\"] = [coord for point in merged_polygon.exterior.coords for coord in point]\n elif isinstance(merged_polygon, MultiPolygon):\n largest_polygon = max(merged_polygon.geoms, key=lambda p: p.area)\n new_annotation[\"segmentation\"] = [coord for point in largest_polygon.exterior.coords for coord in point]\n \n # Ask user about keeping original annotations\n msg_box = QMessageBox(self)\n msg_box.setWindowTitle(\"Merge Annotations\")\n msg_box.setText(\"Do you want to keep the original annotations?\")\n msg_box.setIcon(QMessageBox.Question)\n \n keep_button = msg_box.addButton(\"Keep\", QMessageBox.YesRole)\n delete_button = msg_box.addButton(\"Delete\", QMessageBox.NoRole)\n cancel_button = msg_box.addButton(\"Cancel\", QMessageBox.RejectRole)\n \n msg_box.setDefaultButton(cancel_button)\n msg_box.setEscapeButton(cancel_button)\n \n msg_box.exec_()\n \n if msg_box.clickedButton() == cancel_button:\n return\n \n if msg_box.clickedButton() == delete_button:\n for annotation in original_annotations:\n if annotation in self.image_label.annotations[class_name]:\n self.image_label.annotations[class_name].remove(annotation)\n \n self.image_label.annotations.setdefault(class_name, []).append(new_annotation)\n \n current_name = self.current_slice or self.image_file_name\n self.all_annotations[current_name] = self.image_label.annotations\n \n self.renumber_annotations()\n self.update_annotation_list()\n self.save_current_annotations()\n self.update_slice_list_colors()\n self.image_label.update()\n \n QMessageBox.information(self, \"Merge Complete\", \"Annotations have been merged successfully.\")\n self.auto_save() # Auto-save after merging annotations\n \n \n \n def delete_selected_image(self):\n current_item = self.image_list.currentItem()\n if current_item:\n file_name = current_item.text()\n reply = QMessageBox.question(self, 'Delete Image',\n f\"Are you sure you want to delete the image '{file_name}'?\\n\\n\"\n \"This will remove the image and all its associated annotations.\",\n QMessageBox.Yes | QMessageBox.No, QMessageBox.No)\n \n if reply == QMessageBox.Yes:\n # Remove from all data structures\n self.image_list.takeItem(self.image_list.row(current_item))\n self.image_paths.pop(file_name, None)\n self.all_images = [img for img in self.all_images if img[\"file_name\"] != file_name]\n \n # Remove annotations\n self.all_annotations.pop(file_name, None)\n \n # Handle multi-dimensional images\n base_name = os.path.splitext(file_name)[0]\n if base_name in self.image_slices:\n # Remove slices\n for slice_name, _ in self.image_slices[base_name]:\n self.all_annotations.pop(slice_name, None)\n del self.image_slices[base_name]\n \n # Clear slice list\n self.slice_list.clear()\n \n # Clear current image and slice if it was the removed image\n if self.image_file_name == file_name:\n self.current_image = None\n self.image_file_name = \"\"\n self.current_slice = None\n self.image_label.clear()\n self.annotation_list.clear()\n \n # Switch to another image if available\n if self.image_list.count() > 0:\n next_item = self.image_list.item(0)\n self.image_list.setCurrentItem(next_item)\n self.switch_image(next_item)\n else:\n # No images left\n self.current_image = None\n self.image_file_name = \"\"\n self.current_slice = None\n self.image_label.clear()\n self.annotation_list.clear()\n self.slice_list.clear()\n \n # Update UI\n self.update_ui()\n \n QMessageBox.information(self, \"Image Deleted\", f\"The image '{file_name}' has been deleted.\")\n \n\n def display_image(self):\n if self.current_image:\n if isinstance(self.current_image, QImage):\n pixmap = QPixmap.fromImage(self.current_image)\n elif isinstance(self.current_image, QPixmap):\n pixmap = self.current_image\n else:\n print(f\"Unexpected image type: {type(self.current_image)}\")\n return\n \n if not pixmap.isNull():\n self.image_label.setPixmap(pixmap)\n self.image_label.adjustSize()\n else:\n print(\"Error: Null pixmap\")\n else:\n self.image_label.clear()\n print(\"No current image to display\")\n \n def update_ui(self):\n self.update_image_list()\n self.update_slice_list()\n self.update_class_list()\n self.update_annotation_list()\n self.image_label.update()\n self.update_image_info()\n \n\n def add_class(self, class_name=None, color=None):\n if not self.image_label.check_unsaved_changes(): \n return\n \n if class_name is None:\n while True:\n class_name, ok = QInputDialog.getText(self, \"Add Class\", \"Enter class name:\")\n if not ok:\n print(\"Class addition cancelled\")\n return\n if not class_name.strip():\n QMessageBox.warning(self, \"Invalid Input\", \"Please enter a class name or press Cancel.\")\n continue\n if class_name in self.class_mapping:\n QMessageBox.warning(self, \"Duplicate Class\", f\"The class '{class_name}' already exists. Please choose a different name.\")\n continue\n break\n else:\n # For programmatic addition (e.g., from YOLO predictions)\n if class_name in self.class_mapping:\n print(f\"Class '{class_name}' already exists. Skipping addition.\")\n return\n \n if not isinstance(class_name, str):\n print(f\"Warning: class_name is not a string. Converting {class_name} to string.\")\n class_name = str(class_name)\n \n if color is None:\n color = QColor(Qt.GlobalColor(len(self.image_label.class_colors) % 16 + 7))\n elif isinstance(color, str):\n color = QColor(color)\n \n print(f\"Adding class: {class_name}, color: {color.name()}\")\n \n self.image_label.class_colors[class_name] = color\n self.class_mapping[class_name] = len(self.class_mapping) + 1\n \n try:\n item = QListWidgetItem(class_name)\n \n # Create a color indicator\n pixmap = QPixmap(16, 16)\n pixmap.fill(color)\n item.setIcon(QIcon(pixmap))\n \n # Set visibility state\n item.setData(Qt.UserRole, True)\n \n # Set checkbox\n item.setFlags(item.flags() | Qt.ItemIsUserCheckable)\n item.setCheckState(Qt.Checked)\n \n self.class_list.addItem(item)\n \n self.class_list.setCurrentItem(item)\n self.current_class = class_name\n print(f\"Class added successfully: {class_name}\")\n \n if not self.is_loading_project:\n self.auto_save()\n except Exception as e:\n print(f\"Error adding class: {e}\")\n import traceback\n traceback.print_exc()\n\n if self.class_list.count() > 0:\n self.allButton.setEnabled(True)\n self.clrButton.setEnabled(True)\n \n def update_class_item_color(self, item, color):\n pixmap = QPixmap(16, 16)\n pixmap.fill(color)\n item.setIcon(QIcon(pixmap))\n \n \n \n \n def update_class_list(self):\n self.class_list.clear()\n for class_name, color in self.image_label.class_colors.items():\n item = QListWidgetItem(class_name)\n \n # Create a color indicator\n pixmap = QPixmap(16, 16)\n pixmap.fill(color)\n item.setIcon(QIcon(pixmap))\n \n # Store the visibility state\n item.setData(Qt.UserRole, self.image_label.class_visibility.get(class_name, True))\n \n # Set checkbox\n item.setFlags(item.flags() | Qt.ItemIsUserCheckable)\n item.setCheckState(Qt.Checked if item.data(Qt.UserRole) else Qt.Unchecked)\n \n self.class_list.addItem(item)\n \n # Re-select the current class if it exists\n if self.current_class:\n items = self.class_list.findItems(self.current_class, Qt.MatchExactly)\n if items:\n self.class_list.setCurrentItem(items[0])\n elif self.class_list.count() > 0:\n # If no class is selected, select the first one\n self.class_list.setCurrentItem(self.class_list.item(0))\n\n if self.class_list.count() > 0:\n self.allButton.setEnabled(True)\n self.clrButton.setEnabled(True)\n \n print(f\"Updated class list with {self.class_list.count()} items\")\n \n def update_class_selection(self):\n for i in range(self.class_list.count()):\n item = self.class_list.item(i)\n if item.text() == self.current_class:\n item.setSelected(True)\n else:\n item.setSelected(False)\n \n\n def toggle_class_visibility(self, item):\n class_name = item.text()\n is_visible = item.checkState() == Qt.Checked\n self.image_label.set_class_visibility(class_name, is_visible)\n item.setData(Qt.UserRole, is_visible)\n self.image_label.update()\n \n\n def toggle_all_class(self, checked):\n for i in range(self.class_list.count()):\n item = self.class_list.item(i)\n item.setCheckState(checked)\n # Update image annotations\n self.image_label.update()\n \n \n def change_annotation_class(self):\n selected_items = self.annotation_list.selectedItems()\n if not selected_items:\n QMessageBox.warning(self, \"No Selection\", \"Please select one or more annotations to change class.\")\n return\n \n class_dialog = QDialog(self)\n class_dialog.setWindowTitle(\"Change Class\")\n layout = QVBoxLayout(class_dialog)\n \n class_combo = QComboBox()\n for class_name in self.class_mapping.keys():\n class_combo.addItem(class_name)\n layout.addWidget(class_combo)\n \n button_box = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel)\n button_box.accepted.connect(class_dialog.accept)\n button_box.rejected.connect(class_dialog.reject)\n layout.addWidget(button_box)\n \n if class_dialog.exec_() == QDialog.Accepted:\n new_class = class_combo.currentText()\n current_name = self.current_slice or self.image_file_name\n \n # Get the current maximum number for the new class\n max_number = max([ann.get('number', 0) for ann in self.image_label.annotations.get(new_class, [])] + [0])\n \n for item in selected_items:\n annotation = item.data(Qt.UserRole)\n old_class = annotation['category_name']\n \n # Remove from old class\n self.image_label.annotations[old_class].remove(annotation)\n if not self.image_label.annotations[old_class]:\n del self.image_label.annotations[old_class]\n \n # Add to new class with updated number\n annotation['category_name'] = new_class\n annotation['category_id'] = self.class_mapping[new_class]\n max_number += 1\n annotation['number'] = max_number\n if new_class not in self.image_label.annotations:\n self.image_label.annotations[new_class] = []\n self.image_label.annotations[new_class].append(annotation)\n \n # Update all_annotations\n self.all_annotations[current_name] = self.image_label.annotations\n \n # Renumber all annotations for consistency\n self.renumber_annotations()\n \n self.update_annotation_list()\n self.image_label.update()\n self.save_current_annotations()\n self.update_slice_list_colors()\n self.auto_save()\n \n QMessageBox.information(self, \"Class Changed\", f\"Selected annotations have been changed to class '{new_class}'.\")\n\n\n \n def toggle_tool(self):\n if not self.image_label.check_unsaved_changes():\n return\n \n sender = self.sender()\n if sender is None:\n sender = self.sam_magic_wand_button\n \n if not self.current_class:\n QMessageBox.warning(self, \"No Class Selected\", \"Please select a class before using annotation tools.\")\n sender.setChecked(False)\n return\n \n if self.current_class and self.current_class.startswith(\"Temp-\"):\n QMessageBox.warning(self, \"Invalid Selection\", \"Cannot use annotation tools with temporary classes.\")\n sender.setChecked(False)\n return\n \n other_buttons = [btn for btn in self.tool_group.buttons() if btn != sender]\n \n # Deactivate SAM if we're switching to a different tool\n if sender != self.sam_magic_wand_button and self.image_label.sam_magic_wand_active:\n self.deactivate_sam_magic_wand()\n \n if sender.isChecked():\n # Uncheck all other buttons\n for btn in other_buttons:\n btn.setChecked(False)\n \n # Set the current tool based on the checked button\n if sender == self.polygon_button:\n self.image_label.current_tool = \"polygon\"\n elif sender == self.rectangle_button:\n self.image_label.current_tool = \"rectangle\"\n elif sender == self.sam_magic_wand_button:\n self.image_label.current_tool = \"sam_magic_wand\"\n self.activate_sam_magic_wand()\n elif sender == self.paint_brush_button:\n self.image_label.current_tool = \"paint_brush\"\n self.image_label.setFocus() # Set focus on the image label\n elif sender == self.eraser_button:\n self.image_label.current_tool = \"eraser\"\n self.image_label.setFocus() # Set focus on the image label\n else:\n self.image_label.current_tool = None\n if sender == self.sam_magic_wand_button:\n self.deactivate_sam_magic_wand()\n \n # Update UI based on the current tool\n self.update_ui_for_current_tool()\n\n def wheelEvent(self, event):\n if event.modifiers() == Qt.ControlModifier:\n delta = event.angleDelta().y()\n if self.image_label.current_tool == \"paint_brush\":\n self.paint_brush_size = max(1, self.paint_brush_size + delta // 120)\n print(f\"Paint brush size: {self.paint_brush_size}\")\n elif self.image_label.current_tool == \"eraser\":\n self.eraser_size = max(1, self.eraser_size + delta // 120)\n print(f\"Eraser size: {self.eraser_size}\")\n else:\n super().wheelEvent(event)\n\n\n def update_ui_for_current_tool(self):\n # Disable finish_polygon_button if it still exists in your code\n if hasattr(self, 'finish_polygon_button'):\n self.finish_polygon_button.setEnabled(self.image_label.current_tool in [\"polygon\", \"rectangle\"])\n \n # Update button states\n self.polygon_button.setChecked(self.image_label.current_tool == \"polygon\")\n self.rectangle_button.setChecked(self.image_label.current_tool == \"rectangle\")\n self.sam_magic_wand_button.setChecked(self.image_label.current_tool == \"sam_magic_wand\")\n \n # Enable/disable SAM button based on model availability\n self.sam_magic_wand_button.setEnabled(self.current_sam_model is not None)\n \n # Disable all tools if no class is selected\n tools_enabled = self.current_class is not None and not self.current_class.startswith(\"Temp-\")\n for button in self.tool_group.buttons():\n button.setEnabled(tools_enabled)\n \n # Update cursor based on the current tool\n if self.image_label.current_tool == \"sam_magic_wand\" and self.sam_magic_wand_button.isEnabled():\n self.image_label.setCursor(Qt.CrossCursor)\n else:\n self.image_label.setCursor(Qt.ArrowCursor) \n\n def on_class_selected(self, current=None, previous=None):\n if not self.image_label.check_unsaved_changes():\n return\n \n if current is None:\n current = self.class_list.currentItem()\n \n if current:\n self.current_class = current.text()\n print(f\"Class selected: {self.current_class}\")\n \n if self.current_class.startswith(\"Temp-\"):\n self.disable_annotation_tools()\n else:\n self.enable_annotation_tools()\n else:\n self.current_class = None\n self.disable_annotation_tools()\n \n def disable_annotation_tools(self):\n for button in self.tool_group.buttons():\n button.setChecked(False)\n button.setEnabled(False)\n self.image_label.current_tool = None\n \n def enable_annotation_tools(self):\n for button in self.tool_group.buttons():\n button.setEnabled(True)\n\n def show_class_context_menu(self, position):\n menu = QMenu()\n rename_action = menu.addAction(\"Rename Class\")\n change_color_action = menu.addAction(\"Change Color\")\n delete_action = menu.addAction(\"Delete Class\")\n \n item = self.class_list.itemAt(position)\n if item:\n action = menu.exec_(self.class_list.mapToGlobal(position))\n \n if action == rename_action:\n self.rename_class(item)\n elif action == change_color_action:\n self.change_class_color(item)\n elif action == delete_action:\n self.delete_class(item)\n else:\n QMessageBox.warning(self, \"No Selection\", \"Please select a class to perform actions.\")\n \n def change_class_color(self, item):\n class_name = item.text()\n current_color = self.image_label.class_colors.get(class_name, QColor(Qt.white))\n color = QColorDialog.getColor(current_color, self, f\"Select Color for {class_name}\")\n \n if color.isValid():\n self.image_label.class_colors[class_name] = color\n \n # Update the color indicator\n pixmap = QPixmap(16, 16)\n pixmap.fill(color)\n item.setIcon(QIcon(pixmap))\n \n self.update_annotation_list_colors(class_name, color)\n self.image_label.update()\n self.auto_save() # Auto-save after changing class color\n \n\n \n def rename_class(self, item):\n old_name = item.text()\n new_name, ok = QInputDialog.getText(self, \"Rename Class\", \"Enter new class name:\", text=old_name)\n if ok and new_name and new_name != old_name:\n # Update class mapping\n if old_name in self.class_mapping:\n old_id = self.class_mapping[old_name]\n self.class_mapping[new_name] = old_id\n del self.class_mapping[old_name]\n else:\n print(f\"Warning: Class '{old_name}' not found in class_mapping\")\n return\n \n # Update class colors\n if old_name in self.image_label.class_colors:\n self.image_label.class_colors[new_name] = self.image_label.class_colors.pop(old_name)\n else:\n print(f\"Warning: Class '{old_name}' not found in class_colors\")\n return\n \n # Update annotations for all images and slices\n for image_name, image_annotations in self.all_annotations.items():\n if old_name in image_annotations:\n image_annotations[new_name] = image_annotations.pop(old_name)\n for annotation in image_annotations[new_name]:\n annotation['category_name'] = new_name\n \n # Update current image annotations\n if old_name in self.image_label.annotations:\n self.image_label.annotations[new_name] = self.image_label.annotations.pop(old_name)\n for annotation in self.image_label.annotations[new_name]:\n annotation['category_name'] = new_name\n \n # Update current class if it's the renamed one\n if self.current_class == old_name:\n self.current_class = new_name\n \n # Update annotation list for all images and slices\n self.update_all_annotation_lists()\n \n # Update class list\n item.setText(new_name)\n \n # Update the image label\n self.image_label.update()\n self.auto_save() # Auto-save after renaming a class\n \n print(f\"Class renamed from '{old_name}' to '{new_name}'\")\n \n def delete_class(self, item=None):\n if item is None:\n item = self.class_list.currentItem()\n \n if item is None:\n QMessageBox.warning(self, \"No Selection\", \"Please select a class to delete.\")\n return\n \n class_name = item.text()\n \n # Show confirmation dialog\n reply = QMessageBox.question(self, 'Delete Class',\n f\"Are you sure you want to delete the class '{class_name}'?\\n\\n\"\n \"This will remove all annotations associated with this class.\",\n QMessageBox.Yes | QMessageBox.No, QMessageBox.No)\n \n if reply == QMessageBox.Yes:\n # Proceed with deletion\n # Remove class color\n self.image_label.class_colors.pop(class_name, None)\n \n # Remove class from mapping\n self.class_mapping.pop(class_name, None)\n \n # Remove annotations for this class from all images\n for image_annotations in self.all_annotations.values():\n image_annotations.pop(class_name, None)\n \n # Remove annotations for this class from current image\n self.image_label.annotations.pop(class_name, None)\n \n # Update annotation list\n self.update_annotation_list()\n \n # Remove class from list\n row = self.class_list.row(item)\n self.class_list.takeItem(row)\n \n # Update current_class\n if self.current_class == class_name:\n self.current_class = None\n if self.class_list.count() > 0:\n self.class_list.setCurrentRow(0)\n self.on_class_selected(self.class_list.item(0))\n else:\n self.disable_annotation_tools()\n \n self.image_label.update()\n \n # Inform the user\n QMessageBox.information(self, \"Class Deleted\", f\"The class '{class_name}' has been deleted.\")\n self.auto_save() # Auto-save after deleting a class\n else:\n # User cancelled the operation\n QMessageBox.information(self, \"Deletion Cancelled\", \"The class deletion was cancelled.\")\n \n \n def finish_polygon(self):\n if self.image_label.current_tool == \"polygon\" and len(self.image_label.current_annotation) > 2:\n if self.current_class is None:\n QMessageBox.warning(self, \"No Class Selected\", \"Please select a class before finishing the annotation.\")\n return\n \n # Create a polygon from the current annotation\n polygon = Polygon(self.image_label.current_annotation)\n \n # Define the image boundary as a rectangle\n image_boundary = Polygon([(0, 0), (self.current_image.width(), 0), \n (self.current_image.width(), self.current_image.height()), \n (0, self.current_image.height())])\n \n # Intersect the polygon with the image boundary\n clipped_polygon = polygon.intersection(image_boundary)\n \n if clipped_polygon.is_empty:\n QMessageBox.warning(self, \"Invalid Annotation\", \"The annotation is completely outside the image boundaries.\")\n self.image_label.clear_current_annotation()\n self.image_label.update()\n return\n \n # Convert the clipped polygon to a segmentation format\n if isinstance(clipped_polygon, Polygon):\n segmentation = [coord for point in clipped_polygon.exterior.coords for coord in point]\n elif isinstance(clipped_polygon, MultiPolygon):\n largest_polygon = max(clipped_polygon.geoms, key=lambda p: p.area)\n segmentation = [coord for point in largest_polygon.exterior.coords for coord in point]\n else:\n QMessageBox.warning(self, \"Invalid Annotation\", \"The annotation could not be processed.\")\n return\n \n new_annotation = {\n \"segmentation\": segmentation,\n \"category_id\": self.class_mapping[self.current_class],\n \"category_name\": self.current_class,\n }\n self.image_label.annotations.setdefault(self.current_class, []).append(new_annotation)\n self.add_annotation_to_list(new_annotation)\n self.image_label.clear_current_annotation()\n self.image_label.drawing_polygon = False # Reset the drawing_polygon flag\n self.image_label.reset_annotation_state()\n self.image_label.update()\n \n # Save the current annotations\n self.save_current_annotations()\n \n # Update the slice list colors\n self.update_slice_list_colors()\n self.auto_save() # Auto-save after adding a polygon annotation\n\n\n def highlight_annotation(self, item):\n self.image_label.highlighted_annotation = item.data(Qt.UserRole)\n self.image_label.update()\n\n def delete_annotation(self):\n current_item = self.annotation_list.currentItem()\n if current_item:\n annotation = current_item.data(Qt.UserRole)\n category_name = annotation['category_name']\n self.image_label.annotations[category_name].remove(annotation)\n self.annotation_list.takeItem(self.annotation_list.row(current_item))\n self.image_label.highlighted_annotation = None\n self.image_label.update()\n\n def add_annotation_to_list(self, annotation):\n class_name = annotation['category_name']\n color = self.image_label.class_colors.get(class_name, QColor(Qt.white))\n annotations = self.image_label.annotations.get(class_name, [])\n number = max([ann.get('number', 0) for ann in annotations] + [0]) + 1\n annotation['number'] = number\n area = calculate_area(annotation)\n item_text = f\"{class_name} - {number:<3} Area: {area:.2f}\"\n \n item = QListWidgetItem(item_text)\n item.setData(Qt.UserRole, annotation)\n item.setForeground(color)\n self.annotation_list.addItem(item)\n \n # Clear the current selection\n self.annotation_list.clearSelection()\n self.image_label.highlighted_annotations.clear()\n self.image_label.update()\n \n \n\n\n def zoom_in(self):\n new_zoom = min(self.image_label.zoom_factor + 0.1, 5.0)\n self.set_zoom(new_zoom)\n\n def zoom_out(self):\n new_zoom = max(self.image_label.zoom_factor - 0.1, 0.1)\n self.set_zoom(new_zoom)\n\n def set_zoom(self, zoom_factor):\n self.image_label.set_zoom(zoom_factor)\n self.zoom_slider.setValue(int(zoom_factor * 100))\n self.image_label.update() \n\n def zoom_image(self):\n zoom_factor = self.zoom_slider.value() / 100\n self.set_zoom(zoom_factor)\n\n def disable_tools(self):\n self.polygon_button.setEnabled(False)\n self.rectangle_button.setEnabled(False)\n #self.finish_polygon_button.setEnabled(False)\n\n def enable_tools(self):\n self.polygon_button.setEnabled(True)\n self.rectangle_button.setEnabled(True)\n\n \n def finish_rectangle(self):\n if self.image_label.current_rectangle:\n x1, y1, x2, y2 = self.image_label.current_rectangle\n \n # Create a rectangle polygon from the annotation\n rectangle = Polygon([(x1, y1), (x2, y1), (x2, y2), (x1, y2)])\n \n # Define the image boundary as a rectangle\n image_boundary = Polygon([(0, 0), (self.current_image.width(), 0), \n (self.current_image.width(), self.current_image.height()), \n (0, self.current_image.height())])\n \n # Intersect the rectangle with the image boundary\n clipped_rectangle = rectangle.intersection(image_boundary)\n \n if clipped_rectangle.is_empty:\n QMessageBox.warning(self, \"Invalid Annotation\", \"The annotation is completely outside the image boundaries.\")\n self.image_label.current_rectangle = None\n self.image_label.update()\n return\n \n # Convert the clipped rectangle to a segmentation format\n if isinstance(clipped_rectangle, Polygon):\n segmentation = [coord for point in clipped_rectangle.exterior.coords for coord in point]\n elif isinstance(clipped_rectangle, MultiPolygon):\n largest_polygon = max(clipped_rectangle.geoms, key=lambda p: p.area)\n segmentation = [coord for point in largest_polygon.exterior.coords for coord in point]\n else:\n QMessageBox.warning(self, \"Invalid Annotation\", \"The annotation could not be processed.\")\n return\n \n new_annotation = {\n \"segmentation\": segmentation,\n \"category_id\": self.class_mapping[self.current_class],\n \"category_name\": self.current_class,\n }\n self.image_label.annotations.setdefault(self.current_class, []).append(new_annotation)\n self.add_annotation_to_list(new_annotation)\n self.image_label.start_point = None\n self.image_label.end_point = None\n self.image_label.current_rectangle = None\n self.image_label.update()\n \n # Save the current annotations\n self.save_current_annotations()\n \n # Update the slice list colors\n self.update_slice_list_colors()\n self.auto_save()\n\n def enter_edit_mode(self, annotation):\n self.editing_mode = True\n self.disable_tools()\n\n QMessageBox.information(self, \"Edit Mode\", \"You are now in edit mode. Click and drag points to move them, Shift+Click to delete points, or click on edges to add new points.\")\n\n def exit_edit_mode(self):\n self.editing_mode = False\n self.enable_tools()\n\n self.image_label.editing_polygon = None\n self.image_label.editing_point_index = None\n self.image_label.hover_point_index = None\n self.update_annotation_list()\n self.image_label.update()\n\n def highlight_annotation_in_list(self, annotation):\n for i in range(self.annotation_list.count()):\n item = self.annotation_list.item(i)\n if item.data(Qt.UserRole) == annotation:\n self.annotation_list.setCurrentItem(item)\n break\n\n def select_annotation_in_list(self, annotation):\n for i in range(self.annotation_list.count()):\n item = self.annotation_list.item(i)\n if item.data(Qt.UserRole) == annotation:\n self.annotation_list.setCurrentItem(item)\n break\n \n################################################################\n \n def setup_yolo_menu(self):\n yolo_menu = self.menuBar().addMenu(\"&YOLO (beta)\")\n \n # Training submenu\n training_submenu = yolo_menu.addMenu(\"Training\")\n \n load_pretrained_action = QAction(\"Load Pre-trained Model\", self)\n load_pretrained_action.triggered.connect(self.load_yolo_model)\n training_submenu.addAction(load_pretrained_action)\n \n prepare_data_action = QAction(\"Prepare YOLO Dataset\", self)\n prepare_data_action.triggered.connect(self.prepare_yolo_dataset)\n training_submenu.addAction(prepare_data_action)\n \n load_yaml_action = QAction(\"Load Dataset YAML\", self)\n load_yaml_action.triggered.connect(self.load_yolo_yaml)\n training_submenu.addAction(load_yaml_action)\n \n train_action = QAction(\"Train Model\", self)\n train_action.triggered.connect(self.show_train_dialog)\n training_submenu.addAction(train_action)\n \n save_model_action = QAction(\"Save Model\", self)\n save_model_action.triggered.connect(self.save_yolo_model)\n training_submenu.addAction(save_model_action)\n \n # Prediction Settings submenu\n prediction_submenu = yolo_menu.addMenu(\"Prediction Settings\")\n \n load_model_action = QAction(\"Load Model\", self)\n load_model_action.triggered.connect(self.load_prediction_model)\n prediction_submenu.addAction(load_model_action)\n \n set_threshold_action = QAction(\"Set Confidence Threshold\", self)\n set_threshold_action.triggered.connect(self.set_confidence_threshold)\n prediction_submenu.addAction(set_threshold_action)\n \n\n\n \n \n\n def load_yolo_model(self):\n if not hasattr(self, 'current_project_dir'):\n QMessageBox.warning(self, \"No Project\", \"Please open or create a project first.\")\n return\n \n if not self.yolo_trainer:\n self.initialize_yolo_trainer()\n \n if self.yolo_trainer.load_model():\n QMessageBox.information(self, \"Model Loaded\", \"YOLO model loaded successfully.\")\n else:\n QMessageBox.warning(self, \"Load Cancelled\", \"Model loading was cancelled.\")\n \n def prepare_yolo_dataset(self):\n if not hasattr(self, 'current_project_file'):\n QMessageBox.warning(self, \"No Project\", \"Please open or create a project first.\")\n return\n \n if not self.yolo_trainer:\n self.initialize_yolo_trainer()\n \n try:\n yaml_path = self.yolo_trainer.prepare_dataset()\n QMessageBox.information(self, \"Dataset Prepared\", f\"YOLO dataset prepared successfully. YAML file: {yaml_path}\")\n except Exception as e:\n QMessageBox.critical(self, \"Error\", f\"An error occurred while preparing the dataset: {str(e)}\")\n\n\n def load_yolo_yaml(self):\n if not hasattr(self, 'current_project_file'):\n QMessageBox.warning(self, \"No Project\", \"Please open or create a project first.\")\n return\n \n if not self.yolo_trainer:\n self.initialize_yolo_trainer()\n \n try:\n if self.yolo_trainer.load_yaml():\n QMessageBox.information(self, \"YAML Loaded\", \"Dataset YAML loaded successfully.\")\n else:\n QMessageBox.warning(self, \"Load Cancelled\", \"YAML loading was cancelled.\")\n except Exception as e:\n QMessageBox.critical(self, \"Error\", f\"An error occurred while loading the YAML file: {str(e)}\")\n \n\n def save_yolo_model(self):\n if not hasattr(self, 'current_project_file'):\n QMessageBox.warning(self, \"No Project\", \"Please open or create a project first.\")\n return\n \n if not self.yolo_trainer or not self.yolo_trainer.model:\n QMessageBox.warning(self, \"No Model\", \"Please train or load a YOLO model first.\")\n return\n \n try:\n if self.yolo_trainer.save_model():\n QMessageBox.information(self, \"Model Saved\", \"YOLO model saved successfully.\")\n else:\n QMessageBox.warning(self, \"Save Cancelled\", \"Model saving was cancelled.\")\n except Exception as e:\n QMessageBox.critical(self, \"Error\", f\"An error occurred while saving the model: {str(e)}\")\n\n def load_prediction_model(self):\n if not hasattr(self, 'current_project_file'):\n QMessageBox.warning(self, \"No Project\", \"Please open or create a project first.\")\n return\n \n if not self.yolo_trainer:\n self.initialize_yolo_trainer()\n \n dialog = LoadPredictionModelDialog(self)\n if dialog.exec_() == QDialog.Accepted:\n model_path = dialog.model_path\n yaml_path = dialog.yaml_path\n if model_path and yaml_path:\n try:\n result, message = self.yolo_trainer.load_prediction_model(model_path, yaml_path)\n if result:\n QMessageBox.information(self, \"Model Loaded\", \"YOLO model and YAML file loaded successfully for prediction.\")\n if message:\n QMessageBox.warning(self, \"Class Mismatch Warning\", message)\n else:\n QMessageBox.critical(self, \"Error Loading Model\", f\"Could not load the model or YAML file: {message}\")\n except Exception as e:\n QMessageBox.critical(self, \"Error\", f\"An error occurred: {str(e)}\")\n else:\n QMessageBox.warning(self, \"Files Required\", \"Both model and YAML files are required for prediction.\")\n\n def show_train_dialog(self):\n if not self.yolo_trainer:\n QMessageBox.warning(self, \"No Project\", \"Please open or create a project first.\")\n return\n if not self.yolo_trainer.model:\n QMessageBox.warning(self, \"No Model\", \"Please load a pre-trained model first.\")\n return\n if not self.yolo_trainer.yaml_path:\n QMessageBox.warning(self, \"No Dataset\", \"Please prepare or load a dataset YAML first.\")\n return\n\n dialog = QDialog(self)\n dialog.setWindowTitle(\"Train YOLO Model\")\n layout = QVBoxLayout()\n\n epochs_label = QLabel(\"Number of Epochs:\")\n epochs_input = QLineEdit(\"100\")\n layout.addWidget(epochs_label)\n layout.addWidget(epochs_input)\n\n imgsz_label = QLabel(\"Image Size:\")\n imgsz_input = QLineEdit(\"640\")\n layout.addWidget(imgsz_label)\n layout.addWidget(imgsz_input)\n\n button_box = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel)\n button_box.accepted.connect(dialog.accept)\n button_box.rejected.connect(dialog.reject)\n layout.addWidget(button_box)\n\n dialog.setLayout(layout)\n\n if dialog.exec_() == QDialog.Accepted:\n epochs = int(epochs_input.text())\n imgsz = int(imgsz_input.text())\n self.start_training(epochs, imgsz)\n\n def initialize_yolo_trainer(self):\n if hasattr(self, 'current_project_dir'):\n self.yolo_trainer = YOLOTrainer(self.current_project_dir, self)\n else:\n QMessageBox.warning(self, \"No Project\", \"Please open or create a project first.\")\n\n def start_training(self, epochs, imgsz):\n if not hasattr(self, 'training_dialog'):\n self.training_dialog = TrainingInfoDialog(self)\n self.training_dialog.show()\n \n self.yolo_trainer.progress_signal.connect(self.training_dialog.update_info)\n self.yolo_trainer.set_progress_callback(self.training_dialog.update_info)\n self.training_dialog.stop_signal.connect(self.yolo_trainer.stop_training_signal)\n \n self.training_thread = TrainingThread(self.yolo_trainer, epochs, imgsz)\n self.training_thread.finished.connect(self.training_finished)\n self.training_thread.start()\n\n def training_finished(self, results):\n self.training_dialog.stop_button.setEnabled(True)\n self.training_dialog.stop_button.setText(\"Stop Training\")\n self.yolo_trainer.progress_signal.disconnect(self.training_dialog.update_info)\n self.training_dialog.stop_signal.disconnect(self.yolo_trainer.stop_training_signal)\n\n if isinstance(results, str):\n QMessageBox.critical(self, \"Training Error\", f\"An error occurred during training: {results}\")\n else:\n QMessageBox.information(self, \"Training Complete\", \"YOLO model training completed successfully.\")\n\n def set_confidence_threshold(self):\n if not hasattr(self, 'current_project_file'):\n QMessageBox.warning(self, \"No Project\", \"Please open or create a project first.\")\n return\n \n if not self.yolo_trainer:\n self.initialize_yolo_trainer()\n \n current_threshold = self.yolo_trainer.conf_threshold\n new_threshold, ok = QInputDialog.getDouble(self, \"Set Confidence Threshold\", \n \"Enter confidence threshold (0-1):\", \n current_threshold, 0, 1, 2)\n if ok:\n self.yolo_trainer.set_conf_threshold(new_threshold)\n QMessageBox.information(self, \"Threshold Updated\", f\"Confidence threshold set to {new_threshold}\")\n\n def show_predict_dialog(self):\n if not self.yolo_trainer or not self.yolo_trainer.model:\n QMessageBox.warning(self, \"No Model\", \"Please load a YOLO model first.\")\n return\n \n dialog = QDialog(self)\n dialog.setWindowTitle(\"Predict with YOLO Model\")\n layout = QVBoxLayout()\n \n image_list = QListWidget()\n for image_name in self.image_paths.keys():\n image_list.addItem(image_name)\n layout.addWidget(QLabel(\"Select images for prediction:\"))\n layout.addWidget(image_list)\n \n conf_label = QLabel(\"Confidence Threshold:\")\n conf_input = QDoubleSpinBox()\n conf_input.setRange(0, 1)\n conf_input.setSingleStep(0.01)\n conf_input.setValue(self.yolo_trainer.conf_threshold)\n layout.addWidget(conf_label)\n layout.addWidget(conf_input)\n \n button_box = QDialogButtonBox(QDialogButtonBox.Cancel)\n predict_button = QPushButton(\"Predict\")\n button_box.addButton(predict_button, QDialogButtonBox.AcceptRole)\n button_box.accepted.connect(dialog.accept)\n button_box.rejected.connect(dialog.reject)\n layout.addWidget(button_box)\n \n dialog.setLayout(layout)\n \n if dialog.exec_() == QDialog.Accepted:\n selected_images = [item.text() for item in image_list.selectedItems()]\n conf = conf_input.value()\n self.yolo_trainer.set_conf_threshold(conf)\n self.run_predictions(selected_images)\n\n def run_predictions(self, selected_images):\n for image_name in selected_images:\n image_path = self.image_paths[image_name]\n results = self.yolo_trainer.predict(image_path)\n self.process_yolo_results(results, image_name)\n\n def process_yolo_results(self, results, image_name):\n image_path = self.image_paths[image_name]\n image = cv2.imread(image_path)\n if image is None:\n QMessageBox.warning(self, \"Error\", f\"Failed to load image: {image_name}\")\n return\n original_height, original_width = image.shape[:2]\n \n temp_annotations = {}\n \n try:\n results, input_size, original_size = results # Unpack the results, input size, and original size\n input_height, input_width = input_size\n orig_height, orig_width = original_size\n \n scale_x = original_width / orig_width\n scale_y = original_height / orig_height\n \n for result in results:\n boxes = result.boxes\n masks = result.masks\n \n if masks is None:\n print(f\"No masks found for {image_name}\")\n continue\n \n for mask, box in zip(masks, boxes):\n try:\n class_id = int(box.cls)\n class_name = self.yolo_trainer.class_names[class_id]\n score = float(box.conf)\n \n mask_array = mask.data.cpu().numpy()[0]\n # Resize mask to original image size\n mask_array = cv2.resize(mask_array, (orig_width, orig_height))\n contours, _ = cv2.findContours((mask_array > 0.5).astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)\n \n if contours:\n epsilon = 0.005 * cv2.arcLength(contours[0], True)\n approx = cv2.approxPolyDP(contours[0], epsilon, True)\n polygon = approx.flatten().tolist()\n \n # Scale the polygon coordinates\n scaled_polygon = []\n for i in range(0, len(polygon), 2):\n x = polygon[i] * scale_x\n y = polygon[i+1] * scale_y\n scaled_polygon.extend([x, y])\n \n temp_class_name = f\"Temp-{class_name}\"\n if temp_class_name not in temp_annotations:\n temp_annotations[temp_class_name] = []\n \n temp_annotation = {\n \"segmentation\": scaled_polygon,\n \"category_name\": temp_class_name,\n \"score\": score,\n \"temp\": True\n }\n temp_annotations[temp_class_name].append(temp_annotation)\n except IndexError:\n QMessageBox.warning(self, \"Class Mismatch\", \n \"There is a mismatch between the model and the YAML file classes. \"\n \"Please check that the YAML file corresponds to the loaded model.\")\n return\n \n except Exception as e:\n QMessageBox.warning(self, \"Prediction Error\", \n f\"An error occurred during prediction: {str(e)}\\n\\n\"\n \"This might be due to a mismatch between the model and the YAML file classes. \"\n \"Please check that the YAML file corresponds to the loaded model.\")\n return\n \n self.add_temp_classes(temp_annotations)\n self.update_class_list()\n self.image_label.update()\n \n if temp_annotations:\n total_predictions = sum(len(anns) for anns in temp_annotations.values())\n QMessageBox.information(self, \"Review Predictions\", \n f\"Found {total_predictions} predictions for {len(temp_annotations)} classes.\\n\"\n \"Use class visibility checkboxes to review.\\n\"\n \"Press Enter to accept or Esc to reject visible predictions.\")\n else:\n QMessageBox.information(self, \"No Predictions\", \n \"No predictions were found for this image.\")\n \n # Deactivate SAM tool\n self.deactivate_sam_magic_wand()\n \n \n \n def add_temp_classes(self, temp_annotations):\n for temp_class_name, annotations in temp_annotations.items():\n if temp_class_name not in self.image_label.class_colors:\n color = QColor(Qt.GlobalColor(len(self.image_label.class_colors) % 16 + 7))\n self.image_label.class_colors[temp_class_name] = color\n self.image_label.annotations[temp_class_name] = annotations\n \n self.update_class_list()\n \n def verify_current_class(self):\n if self.current_class is None or self.current_class not in self.class_mapping:\n if self.class_list.count() > 0:\n self.class_list.setCurrentRow(0)\n self.on_class_selected(self.class_list.item(0))\n else:\n self.current_class = None\n self.disable_annotation_tools()\n \n def accept_visible_temp_classes(self):\n visible_temp_classes = [item.text() for item in self.class_list.findItems(\"Temp-*\", Qt.MatchWildcard) \n if item.checkState() == Qt.Checked]\n \n for temp_class_name in visible_temp_classes:\n permanent_class_name = temp_class_name[5:] # Remove \"Temp-\" prefix\n if permanent_class_name not in self.image_label.annotations:\n self.add_class(permanent_class_name, self.image_label.class_colors[temp_class_name])\n \n # Get the current maximum number for this class\n current_max = max([ann.get('number', 0) for ann in self.image_label.annotations.get(permanent_class_name, [])] + [0])\n \n for annotation in self.image_label.annotations[temp_class_name]:\n current_max += 1\n annotation['category_name'] = permanent_class_name\n annotation['number'] = current_max\n self.image_label.annotations.setdefault(permanent_class_name, []).append(annotation)\n \n del self.image_label.annotations[temp_class_name]\n del self.image_label.class_colors[temp_class_name]\n \n self.update_class_list()\n current_name = self.current_slice or self.image_file_name\n self.all_annotations[current_name] = self.image_label.annotations\n self.update_annotation_list()\n self.image_label.update()\n self.save_current_annotations()\n \n # Select the first primary class\n self.select_first_primary_class()\n self.verify_current_class()\n \n QMessageBox.information(self, \"Annotations Accepted\", \"Temporary annotations have been accepted and added to the permanent classes.\")\n \n def select_first_primary_class(self):\n for i in range(self.class_list.count()):\n item = self.class_list.item(i)\n if not item.text().startswith(\"Temp-\"):\n self.class_list.setCurrentItem(item)\n self.on_class_selected(item)\n break\n \n def reject_visible_temp_classes(self):\n visible_temp_classes = [item.text() for item in self.class_list.findItems(\"Temp-*\", Qt.MatchWildcard) \n if item.checkState() == Qt.Checked]\n \n for temp_class_name in visible_temp_classes:\n if temp_class_name in self.image_label.annotations:\n del self.image_label.annotations[temp_class_name]\n if temp_class_name in self.image_label.class_colors:\n del self.image_label.class_colors[temp_class_name]\n \n self.update_class_list()\n self.image_label.update()\n \n def is_class_visible(self, class_name):\n items = self.class_list.findItems(class_name, Qt.MatchExactly)\n if items:\n return items[0].checkState() == Qt.Checked\n return False\n \n def check_temp_annotations(self):\n temp_classes = [class_name for class_name in self.image_label.annotations.keys() if class_name.startswith(\"Temp-\")]\n if temp_classes:\n reply = QMessageBox.question(self, 'Temporary Annotations',\n \"There are temporary annotations that will be discarded. Do you want to continue?\",\n QMessageBox.Yes | QMessageBox.No, QMessageBox.No)\n if reply == QMessageBox.Yes:\n for temp_class in temp_classes:\n del self.image_label.annotations[temp_class]\n del self.image_label.class_colors[temp_class]\n self.update_class_list()\n self.update_annotation_list()\n return True\n return False\n return True\n\n def remove_all_temp_annotations(self):\n for image_name in list(self.all_annotations.keys()):\n for class_name in list(self.all_annotations[image_name].keys()):\n if class_name.startswith(\"Temp-\"):\n del self.all_annotations[image_name][class_name]\n if not self.all_annotations[image_name]:\n del self.all_annotations[image_name]\n \n for class_name in list(self.image_label.class_colors.keys()):\n if class_name.startswith(\"Temp-\"):\n del self.image_label.class_colors[class_name]\n \n self.update_class_list()\n self.update_annotation_list()\n self.image_label.update()\n"
},
{
"path": "src/digitalsreeni_image_annotator/coco_json_combiner.py",
"content": "import json\nimport os\nfrom PyQt5.QtWidgets import (QDialog, QVBoxLayout, QHBoxLayout, QPushButton, \n QFileDialog, QLabel, QMessageBox, QApplication)\nfrom PyQt5.QtCore import Qt\n\nclass COCOJSONCombinerDialog(QDialog):\n def __init__(self, parent=None):\n super().__init__(parent)\n self.setWindowTitle(\"COCO JSON Combiner\")\n self.setGeometry(100, 100, 400, 300)\n self.setWindowFlags(self.windowFlags() | Qt.Window)\n self.setWindowModality(Qt.ApplicationModal)\n self.json_files = []\n self.initUI()\n\n def initUI(self):\n layout = QVBoxLayout()\n\n self.file_labels = []\n for i in range(5):\n file_layout = QHBoxLayout()\n label = QLabel(f\"File {i+1}: Not selected\")\n self.file_labels.append(label)\n file_layout.addWidget(label)\n select_button = QPushButton(f\"Select File {i+1}\")\n select_button.clicked.connect(lambda checked, x=i: self.select_file(x))\n file_layout.addWidget(select_button)\n layout.addLayout(file_layout)\n\n self.combine_button = QPushButton(\"Combine JSON Files\")\n self.combine_button.clicked.connect(self.combine_json_files)\n self.combine_button.setEnabled(False)\n layout.addWidget(self.combine_button)\n\n self.setLayout(layout)\n\n def select_file(self, index):\n file_name, _ = QFileDialog.getOpenFileName(self, f\"Select COCO JSON File {index+1}\", \"\", \"JSON Files (*.json)\")\n if file_name:\n if file_name not in self.json_files:\n self.json_files.append(file_name)\n self.file_labels[index].setText(f\"File {index+1}: {os.path.basename(file_name)}\")\n self.combine_button.setEnabled(True)\n else:\n QMessageBox.warning(self, \"Duplicate File\", \"This file has already been selected.\")\n QApplication.processEvents()\n\n\n def combine_json_files(self):\n if not self.json_files:\n QMessageBox.warning(self, \"No Files\", \"Please select at least one JSON file to combine.\")\n return\n \n combined_data = {\n \"images\": [],\n \"annotations\": [],\n \"categories\": []\n }\n image_file_names = set()\n next_image_id = 1\n next_annotation_id = 1\n \n try:\n for file_path in self.json_files:\n with open(file_path, 'r') as f:\n data = json.load(f)\n \n # Combine categories\n category_id_map = {}\n for category in data.get('categories', []):\n existing_category = next((c for c in combined_data['categories'] if c['name'] == category['name']), None)\n if existing_category:\n category_id_map[category['id']] = existing_category['id']\n else:\n new_id = len(combined_data['categories']) + 1\n category_id_map[category['id']] = new_id\n category['id'] = new_id\n combined_data['categories'].append(category)\n \n # Combine images and annotations\n image_id_map = {}\n for image in data.get('images', []):\n if image['file_name'] not in image_file_names:\n image_file_names.add(image['file_name'])\n image_id_map[image['id']] = next_image_id\n image['id'] = next_image_id\n combined_data['images'].append(image)\n next_image_id += 1\n \n for annotation in data.get('annotations', []):\n if annotation['image_id'] in image_id_map:\n annotation['id'] = next_annotation_id\n annotation['image_id'] = image_id_map[annotation['image_id']]\n annotation['category_id'] = category_id_map[annotation['category_id']]\n combined_data['annotations'].append(annotation)\n next_annotation_id += 1\n \n output_file, _ = QFileDialog.getSaveFileName(self, \"Save Combined JSON\", \"\", \"JSON Files (*.json)\")\n if output_file:\n with open(output_file, 'w') as f:\n json.dump(combined_data, f, indent=2)\n QMessageBox.information(self, \"Success\", f\"Combined JSON saved to {output_file}\")\n \n except Exception as e:\n QMessageBox.critical(self, \"Error\", f\"An error occurred while combining JSON files: {str(e)}\")\n\n\n\n def show_centered(self, parent):\n parent_geo = parent.geometry()\n self.move(parent_geo.center() - self.rect().center())\n self.show()\n\ndef show_coco_json_combiner(parent):\n dialog = COCOJSONCombinerDialog(parent)\n dialog.show_centered(parent)\n return dialog"
},
{
"path": "src/digitalsreeni_image_annotator/constants.py",
"content": "\"\"\"\nConstants for the Image Annotator application.\n\nThis module contains constant values used across the application.\n\n@DigitalSreeni\nDr. Sreenivas Bhattiprolu\n\"\"\"\n\n# File dialog filters\nIMAGE_FILE_FILTER = \"Image Files (*.png *.jpg *.bmp)\"\nJSON_FILE_FILTER = \"JSON Files (*.json)\"\n\n# Default window size\nDEFAULT_WINDOW_WIDTH = 1400\nDEFAULT_WINDOW_HEIGHT = 800\n\n# Zoom settings\nMIN_ZOOM = 10\nMAX_ZOOM = 500\nDEFAULT_ZOOM = 100\n\n# Annotation settings\nDEFAULT_FILL_OPACITY = 0.3\n\n"
},
{
"path": "src/digitalsreeni_image_annotator/dataset_splitter.py",
"content": "import os\nimport json\nimport shutil\nimport random\nfrom PyQt5.QtWidgets import (QDialog, QVBoxLayout, QHBoxLayout, QPushButton, QFileDialog, \n QLabel, QSpinBox, QRadioButton, QButtonGroup, QMessageBox, QComboBox)\nfrom PyQt5.QtCore import Qt\nimport yaml\nfrom PIL import Image\n\nclass DatasetSplitterTool(QDialog):\n def __init__(self, parent=None):\n super().__init__(parent)\n self.setWindowTitle(\"Dataset Splitter\")\n self.setGeometry(100, 100, 500, 300)\n self.setWindowFlags(self.windowFlags() | Qt.Window)\n self.initUI()\n\n def initUI(self):\n layout = QVBoxLayout()\n\n # Option selection\n options_layout = QVBoxLayout()\n self.images_only_radio = QRadioButton(\"Images Only\")\n options_layout.addWidget(self.images_only_radio)\n\n images_annotations_layout = QHBoxLayout()\n self.images_annotations_radio = QRadioButton(\"Images and Annotations\")\n images_annotations_layout.addWidget(self.images_annotations_radio)\n self.select_json_button = QPushButton(\"Upload COCO JSON File\")\n self.select_json_button.clicked.connect(self.select_json_file)\n self.select_json_button.setEnabled(False)\n images_annotations_layout.addWidget(self.select_json_button)\n options_layout.addLayout(images_annotations_layout)\n\n layout.addLayout(options_layout)\n \n option_group = QButtonGroup(self)\n option_group.addButton(self.images_only_radio)\n option_group.addButton(self.images_annotations_radio)\n \n self.images_only_radio.setChecked(True)\n\n # Percentage inputs\n train_layout = QHBoxLayout()\n train_layout.addWidget(QLabel(\"Train %:\"))\n self.train_percent = QSpinBox()\n self.train_percent.setRange(0, 100)\n self.train_percent.setValue(70)\n train_layout.addWidget(self.train_percent)\n layout.addLayout(train_layout)\n\n val_layout = QHBoxLayout()\n val_layout.addWidget(QLabel(\"Validation %:\"))\n self.val_percent = QSpinBox()\n self.val_percent.setRange(0, 100)\n self.val_percent.setValue(30)\n val_layout.addWidget(self.val_percent)\n layout.addLayout(val_layout)\n\n test_layout = QHBoxLayout()\n test_layout.addWidget(QLabel(\"Test %:\"))\n self.test_percent = QSpinBox()\n self.test_percent.setRange(0, 100)\n self.test_percent.setValue(0)\n test_layout.addWidget(self.test_percent)\n layout.addLayout(test_layout)\n\n # Format selection\n self.format_selection_layout = QHBoxLayout()\n self.format_label = QLabel(\"Output Format:\")\n self.format_combo = QComboBox()\n self.format_combo.addItems([\"COCO JSON\", \"YOLO\"])\n self.format_combo.setEnabled(False)\n self.format_selection_layout.addWidget(self.format_label)\n self.format_selection_layout.addWidget(self.format_combo)\n options_layout.addLayout(self.format_selection_layout)\n\n # Buttons\n self.select_input_button = QPushButton(\"Select Input Directory\")\n self.select_input_button.clicked.connect(self.select_input_directory)\n layout.addWidget(self.select_input_button)\n\n self.select_output_button = QPushButton(\"Select Output Directory\")\n self.select_output_button.clicked.connect(self.select_output_directory)\n layout.addWidget(self.select_output_button)\n\n self.split_button = QPushButton(\"Split Dataset\")\n self.split_button.clicked.connect(self.split_dataset)\n layout.addWidget(self.split_button)\n\n self.setLayout(layout)\n\n self.input_directory = \"\"\n self.output_directory = \"\"\n self.json_file = \"\"\n\n # Connect radio buttons to enable/disable JSON selection\n self.images_only_radio.toggled.connect(self.toggle_json_selection)\n self.images_annotations_radio.toggled.connect(self.toggle_json_selection)\n\n def toggle_json_selection(self):\n is_annotations = self.images_annotations_radio.isChecked()\n self.select_json_button.setEnabled(is_annotations)\n self.format_combo.setEnabled(is_annotations)\n\n def select_input_directory(self):\n self.input_directory = QFileDialog.getExistingDirectory(self, \"Select Input Directory\")\n\n def select_output_directory(self):\n self.output_directory = QFileDialog.getExistingDirectory(self, \"Select Output Directory\")\n\n def select_json_file(self):\n self.json_file, _ = QFileDialog.getOpenFileName(self, \"Select COCO JSON File\", \"\", \"JSON Files (*.json)\")\n\n def split_dataset(self):\n if not self.input_directory or not self.output_directory:\n QMessageBox.warning(self, \"Error\", \"Please select input and output directories.\")\n return\n\n if self.images_annotations_radio.isChecked() and not self.json_file:\n QMessageBox.warning(self, \"Error\", \"Please select a COCO JSON file.\")\n return\n\n train_percent = self.train_percent.value()\n val_percent = self.val_percent.value()\n test_percent = self.test_percent.value()\n\n if train_percent + val_percent + test_percent != 100:\n QMessageBox.warning(self, \"Error\", \"Percentages must add up to 100%.\")\n return\n\n if self.images_only_radio.isChecked():\n self.split_images_only()\n else:\n self.split_images_and_annotations()\n\n def split_images_only(self):\n image_files = [f for f in os.listdir(self.input_directory) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.tif', '.tiff'))]\n random.shuffle(image_files)\n\n train_split = int(len(image_files) * self.train_percent.value() / 100)\n val_split = int(len(image_files) * self.val_percent.value() / 100)\n\n train_images = image_files[:train_split]\n val_images = image_files[train_split:train_split + val_split]\n test_images = image_files[train_split + val_split:]\n\n for subset, images in [(\"train\", train_images), \n (\"val\", val_images), \n (\"test\", test_images)]:\n if images: # Only create directories and copy images if there are images for this split\n subset_dir = os.path.join(self.output_directory, subset)\n os.makedirs(subset_dir, exist_ok=True)\n self.copy_images(images, subset, images_only=True)\n\n QMessageBox.information(self, \"Success\", \"Dataset split successfully!\")\n\n def split_images_and_annotations(self):\n with open(self.json_file, 'r') as f:\n coco_data = json.load(f)\n\n image_files = [img['file_name'] for img in coco_data['images']]\n random.shuffle(image_files)\n\n train_split = int(len(image_files) * self.train_percent.value() / 100)\n val_split = int(len(image_files) * self.val_percent.value() / 100)\n\n train_images = image_files[:train_split]\n val_images = image_files[train_split:train_split + val_split]\n test_images = image_files[train_split + val_split:]\n\n # Create main directories\n os.makedirs(self.output_directory, exist_ok=True)\n \n if self.format_combo.currentText() == \"COCO JSON\":\n self.split_coco_format(coco_data, train_images, val_images, test_images)\n else: # YOLO format\n self.split_yolo_format(coco_data, train_images, val_images, test_images)\n\n def copy_images(self, image_list, subset, images_only=False):\n if not image_list:\n return\n \n if images_only:\n subset_dir = os.path.join(self.output_directory, subset)\n else:\n subset_dir = os.path.join(self.output_directory, subset, \"images\")\n os.makedirs(subset_dir, exist_ok=True)\n \n for image in image_list:\n src = os.path.join(self.input_directory, image)\n dst = os.path.join(subset_dir, image)\n shutil.copy2(src, dst)\n\n def create_subset_annotations(self, coco_data, subset_images):\n subset_images_data = [img for img in coco_data['images'] if img['file_name'] in subset_images]\n subset_image_ids = [img['id'] for img in subset_images_data]\n \n return {\n \"images\": subset_images_data,\n \"annotations\": [ann for ann in coco_data['annotations'] if ann['image_id'] in subset_image_ids],\n \"categories\": coco_data['categories']\n }\n\n def split_coco_format(self, coco_data, train_images, val_images, test_images):\n # Only create directories and save annotations for non-empty splits\n for subset, images in [(\"train\", train_images), \n (\"val\", val_images), \n (\"test\", test_images)]:\n if images: # Only process if there are images in this split\n subset_dir = os.path.join(self.output_directory, subset)\n os.makedirs(subset_dir, exist_ok=True) # Create the subset directory first\n os.makedirs(os.path.join(subset_dir, \"images\"), exist_ok=True)\n self.copy_images(images, subset, images_only=False)\n \n # Create and save annotations for this subset\n subset_data = self.create_subset_annotations(coco_data, images)\n self.save_coco_annotations(subset_data, subset)\n\n QMessageBox.information(self, \"Success\", \"Dataset and COCO annotations split successfully!\")\n\n def save_coco_annotations(self, data, subset):\n subset_dir = os.path.join(self.output_directory, subset)\n os.makedirs(subset_dir, exist_ok=True)\n output_file = os.path.join(subset_dir, f\"{subset}_annotations.json\")\n with open(output_file, 'w') as f:\n json.dump(data, f, indent=2)\n\n def split_yolo_format(self, coco_data, train_images, val_images, test_images):\n # Create directories only for non-empty splits\n yaml_paths = {}\n for subset, images in [(\"train\", train_images), \n (\"val\", val_images), \n (\"test\", test_images)]:\n if images: # Only create directories if there are images for this split\n subset_dir = os.path.join(self.output_directory, subset)\n os.makedirs(os.path.join(subset_dir, \"images\"), exist_ok=True)\n os.makedirs(os.path.join(subset_dir, \"labels\"), exist_ok=True)\n yaml_paths[subset] = f'./{subset}/images'\n\n # Create class mapping (COCO to YOLO indices)\n categories = {cat[\"id\"]: i for i, cat in enumerate(coco_data[\"categories\"])}\n\n # Process each non-empty subset\n for subset, images in [(\"train\", train_images), \n (\"val\", val_images), \n (\"test\", test_images)]:\n if not images: # Skip if no images in this split\n continue\n \n images_dir = os.path.join(self.output_directory, subset, \"images\")\n labels_dir = os.path.join(self.output_directory, subset, \"labels\")\n \n for image_file in images:\n # Copy image\n src = os.path.join(self.input_directory, image_file)\n shutil.copy2(src, os.path.join(images_dir, image_file))\n \n # Get image dimensions\n img = Image.open(src)\n img_width, img_height = img.size\n \n # Get annotations for this image\n image_id = next(img[\"id\"] for img in coco_data[\"images\"] if img[\"file_name\"] == image_file)\n annotations = [ann for ann in coco_data[\"annotations\"] if ann[\"image_id\"] == image_id]\n \n # Create YOLO format labels\n label_file = os.path.join(labels_dir, os.path.splitext(image_file)[0] + \".txt\")\n with open(label_file, \"w\") as f:\n for ann in annotations:\n # Convert COCO class id to YOLO class id\n yolo_class = categories[ann[\"category_id\"]]\n \n # Convert COCO bbox to YOLO format\n x, y, w, h = ann[\"bbox\"]\n x_center = (x + w/2) / img_width\n y_center = (y + h/2) / img_height\n w = w / img_width\n h = h / img_height\n \n f.write(f\"{yolo_class} {x_center:.6f} {y_center:.6f} {w:.6f} {h:.6f}\\n\")\n\n # Create data.yaml with only the relevant paths\n yaml_data = {\n 'nc': len(categories),\n 'names': [cat[\"name\"] for cat in sorted(coco_data[\"categories\"], key=lambda x: categories[x[\"id\"]])]\n }\n yaml_data.update(yaml_paths) # Add only paths for non-empty splits\n\n with open(os.path.join(self.output_directory, 'data.yaml'), 'w') as f:\n yaml.dump(yaml_data, f, default_flow_style=False)\n\n QMessageBox.information(self, \"Success\", \"Dataset and YOLO annotations split successfully!\")\n\n def show_centered(self, parent):\n parent_geo = parent.geometry()\n self.move(parent_geo.center() - self.rect().center())\n self.show()"
},
{
"path": "src/digitalsreeni_image_annotator/default_stylesheet.py",
"content": "default_stylesheet = \"\"\"\nQWidget {\n background-color: #F0F0F0;\n color: #333333;\n font-family: Arial, sans-serif;\n}\n\nQMainWindow {\n background-color: #FFFFFF;\n}\n\nQPushButton {\n background-color: #E0E0E0;\n border: 1px solid #BBBBBB;\n padding: 5px 10px;\n border-radius: 3px;\n color: #333333;\n}\n\nQPushButton:hover {\n background-color: #D0D0D0;\n}\n\nQPushButton:pressed {\n background-color: #C0C0C0;\n}\n\nQPushButton:checked {\n background-color: #A0A0A0;\n border: 2px solid #808080;\n color: #FFFFFF;\n}\n\n\nQListWidget, QTreeWidget {\n background-color: #FFFFFF;\n border: 1px solid #CCCCCC;\n border-radius: 3px;\n}\n\n\nQListWidget::item:selected {\n background-color: #E0E0E0;\n color: #333333;\n}\n\n\nQLabel {\n color: #333333;\n}\n\nQLabel.section-header {\n font-weight: bold;\n font-size: 14px;\n padding: 5px 0;\n color: #333333; /* Dark color for visibility in light mode */\n}\n\n\nQLineEdit, QTextEdit, QPlainTextEdit {\n background-color: #FFFFFF;\n border: 1px solid #CCCCCC;\n color: #333333;\n padding: 2px;\n border-radius: 3px;\n}\n\nQSlider::groove:horizontal {\n background: #CCCCCC;\n height: 8px;\n border-radius: 4px;\n}\n\nQSlider::handle:horizontal {\n background: #888888;\n width: 18px;\n margin-top: -5px;\n margin-bottom: -5px;\n border-radius: 9px;\n}\n\nQSlider::handle:horizontal:hover {\n background: #666666;\n}\n\nQScrollBar:vertical, QScrollBar:horizontal {\n background-color: #F0F0F0;\n width: 12px;\n height: 12px;\n}\n\nQScrollBar::handle:vertical, QScrollBar::handle:horizontal {\n background-color: #CCCCCC;\n border-radius: 6px;\n min-height: 20px;\n}\n\nQScrollBar::handle:vertical:hover, QScrollBar::handle:horizontal:hover {\n background-color: #BBBBBB;\n}\n\nQScrollBar::add-line, QScrollBar::sub-line {\n background: none;\n}\n\nQMenuBar {\n background-color: #F0F0F0;\n}\n\nQMenuBar::item {\n padding: 5px 10px;\n background-color: transparent;\n}\n\nQMenuBar::item:selected {\n background-color: #E0E0E0;\n}\n\nQMenu {\n background-color: #FFFFFF;\n border: 1px solid #CCCCCC;\n}\n\nQMenu::item {\n padding: 5px 20px 5px 20px;\n}\n\nQMenu::item:selected {\n background-color: #E0E0E0;\n}\n\nQToolTip {\n background-color: #FFFFFF;\n color: #333333;\n border: 1px solid #CCCCCC;\n}\n\nQStatusBar {\n background-color: #F0F0F0;\n color: #666666;\n}\n\nQListWidget::item {\n color: none;\n}\n\"\"\""
},
{
"path": "src/digitalsreeni_image_annotator/dicom_converter.py",
"content": "import os\nimport json\nimport numpy as np\nfrom datetime import datetime\nfrom PyQt5.QtWidgets import (QDialog, QVBoxLayout, QHBoxLayout, QPushButton, QFileDialog, \n QLabel, QProgressDialog, QRadioButton, QButtonGroup, \n QMessageBox, QApplication, QGroupBox)\nfrom PyQt5.QtCore import Qt\nimport pydicom\nfrom pydicom.pixel_data_handlers.util import apply_voi_lut\nimport tifffile\n\nclass DicomConverter(QDialog):\n def __init__(self, parent=None):\n super().__init__(parent)\n self.setWindowTitle(\"DICOM to TIFF Converter\")\n self.setGeometry(100, 100, 600, 300)\n self.setWindowFlags(self.windowFlags() | Qt.Window)\n self.setWindowModality(Qt.ApplicationModal) # Add modal behavior\n \n # Initialize variables first\n self.input_file = \"\"\n self.output_directory = \"\"\n \n self.initUI()\n \n def initUI(self):\n layout = QVBoxLayout()\n layout.setSpacing(10) # Add consistent spacing\n\n # File Selection Group\n file_group = QGroupBox(\"File Selection\")\n file_layout = QVBoxLayout()\n\n # Input file selection\n input_layout = QHBoxLayout()\n self.input_label = QLabel(\"No DICOM file selected\")\n self.input_label.setMinimumWidth(100)\n self.input_label.setMaximumWidth(300)\n self.input_label.setWordWrap(True)\n self.select_input_btn = QPushButton(\"Select DICOM File\")\n self.select_input_btn.clicked.connect(self.select_input)\n input_layout.addWidget(self.select_input_btn)\n input_layout.addWidget(self.input_label, 1)\n file_layout.addLayout(input_layout)\n\n # Output directory selection\n output_layout = QHBoxLayout()\n self.output_label = QLabel(\"No output directory selected\")\n self.output_label.setMinimumWidth(100)\n self.output_label.setMaximumWidth(300)\n self.output_label.setWordWrap(True)\n self.select_output_btn = QPushButton(\"Select Output Directory\")\n self.select_output_btn.clicked.connect(self.select_output)\n output_layout.addWidget(self.select_output_btn)\n output_layout.addWidget(self.output_label, 1)\n file_layout.addLayout(output_layout)\n\n file_group.setLayout(file_layout)\n layout.addWidget(file_group)\n\n # Output Format Group\n format_group = QGroupBox(\"Output Format\")\n format_layout = QVBoxLayout()\n \n self.stack_radio = QRadioButton(\"Single TIFF Stack\")\n self.individual_radio = QRadioButton(\"Individual TIFF Files\")\n self.stack_radio.setChecked(True)\n \n format_layout.addWidget(self.stack_radio)\n format_layout.addWidget(self.individual_radio)\n format_group.setLayout(format_layout)\n layout.addWidget(format_group)\n\n # Metadata info\n metadata_group = QGroupBox(\"Metadata Information\")\n metadata_layout = QVBoxLayout()\n metadata_label = QLabel(\"DICOM metadata will be saved as JSON file in the output directory\")\n metadata_label.setStyleSheet(\"color: gray; font-style: italic;\")\n metadata_label.setWordWrap(True)\n metadata_layout.addWidget(metadata_label)\n metadata_group.setLayout(metadata_layout)\n layout.addWidget(metadata_group)\n\n # Convert button\n self.convert_btn = QPushButton(\"Convert\")\n self.convert_btn.clicked.connect(self.convert_dicom)\n layout.addWidget(self.convert_btn)\n\n self.setLayout(layout)\n \n def select_input(self):\n try:\n file_filter = \"DICOM files (*.dcm *.DCM);;All files (*.*)\"\n file_name, _ = QFileDialog.getOpenFileName(\n self,\n \"Select DICOM File\",\n \"\",\n file_filter,\n options=QFileDialog.Options()\n )\n \n if file_name:\n self.input_file = file_name\n self.input_label.setText(self.truncate_path(file_name))\n self.input_label.setToolTip(file_name)\n QApplication.processEvents()\n \n except Exception as e:\n QMessageBox.critical(self, \"Error\", f\"Error selecting input file: {str(e)}\")\n \n def select_output(self):\n try:\n directory = QFileDialog.getExistingDirectory(\n self,\n \"Select Output Directory\",\n \"\",\n QFileDialog.ShowDirsOnly | QFileDialog.DontResolveSymlinks\n )\n \n if directory:\n self.output_directory = directory\n self.output_label.setText(self.truncate_path(directory))\n self.output_label.setToolTip(directory)\n QApplication.processEvents()\n \n except Exception as e:\n QMessageBox.critical(self, \"Error\", f\"Error selecting output directory: {str(e)}\")\n \n def truncate_path(self, path, max_length=40):\n if len(path) <= max_length:\n return path\n \n filename = os.path.basename(path)\n directory = os.path.dirname(path)\n \n if len(filename) > max_length - 5:\n return f\"...{filename[-(max_length-5):]}\"\n \n available_length = max_length - len(filename) - 5\n return f\"...{directory[-available_length:]}{os.sep}{filename}\"\n\n def extract_metadata(self, ds):\n \"\"\"Extract relevant metadata from DICOM dataset.\"\"\"\n metadata = {\n \"PatientID\": getattr(ds, \"PatientID\", \"Unknown\"),\n \"PatientName\": str(getattr(ds, \"PatientName\", \"Unknown\")),\n \"StudyDate\": getattr(ds, \"StudyDate\", \"Unknown\"),\n \"SeriesDescription\": getattr(ds, \"SeriesDescription\", \"Unknown\"),\n \"Modality\": getattr(ds, \"Modality\", \"Unknown\"),\n \"Manufacturer\": getattr(ds, \"Manufacturer\", \"Unknown\"),\n \"InstitutionName\": getattr(ds, \"InstitutionName\", \"Unknown\"),\n \"PixelSpacing\": getattr(ds, \"PixelSpacing\", [1, 1]),\n \"SliceThickness\": getattr(ds, \"SliceThickness\", 1),\n \"ImageOrientation\": getattr(ds, \"ImageOrientationPatient\", [1,0,0,0,1,0]),\n \"ImagePosition\": getattr(ds, \"ImagePositionPatient\", [0,0,0]),\n \"WindowCenter\": getattr(ds, \"WindowCenter\", None),\n \"WindowWidth\": getattr(ds, \"WindowWidth\", None),\n \"RescaleIntercept\": getattr(ds, \"RescaleIntercept\", 0),\n \"RescaleSlope\": getattr(ds, \"RescaleSlope\", 1),\n \"BitsAllocated\": getattr(ds, \"BitsAllocated\", 16),\n \"PixelRepresentation\": getattr(ds, \"PixelRepresentation\", 0),\n \"ConversionDate\": datetime.now().strftime(\"%Y-%m-%d %H:%M:%S\")\n }\n return metadata\n\n def apply_window_level(self, image, ds):\n \"\"\"Apply window/level if present in DICOM.\"\"\"\n try:\n if hasattr(ds, 'WindowCenter') and hasattr(ds, 'WindowWidth'):\n return apply_voi_lut(image, ds)\n except:\n pass\n return image\n\n\n def convert_dicom(self):\n if not self.input_file or not self.output_directory:\n QMessageBox.warning(self, \"Error\", \"Please select both input file and output directory\")\n return\n \n try:\n # Create progress dialog\n progress = QProgressDialog(\"Processing DICOM file...\", \"Cancel\", 0, 100, self)\n progress.setWindowModality(Qt.WindowModal)\n progress.setMinimumWidth(400)\n progress.show()\n \n # Verify DICOM file\n if not pydicom.misc.is_dicom(self.input_file):\n raise ValueError(\"Selected file is not a valid DICOM file\")\n \n # Read DICOM data\n print(\"Reading DICOM file...\")\n progress.setLabelText(\"Reading DICOM file...\")\n progress.setValue(20)\n \n ds = pydicom.dcmread(self.input_file)\n series_metadata = self.extract_metadata(ds)\n \n # Process pixel data\n print(\"Processing pixel data...\")\n progress.setLabelText(\"Processing pixel data...\")\n progress.setValue(40)\n \n pixel_array = ds.pixel_array\n original_dtype = pixel_array.dtype\n print(f\"Original data type: {original_dtype}\")\n print(f\"Original data range: {pixel_array.min()} to {pixel_array.max()}\")\n \n # Apply rescale slope and intercept\n if hasattr(ds, 'RescaleSlope') or hasattr(ds, 'RescaleIntercept'):\n slope = getattr(ds, 'RescaleSlope', 1)\n intercept = getattr(ds, 'RescaleIntercept', 0)\n print(f\"Applying rescale slope ({slope}) and intercept ({intercept})\")\n pixel_array = (pixel_array * slope + intercept)\n \n # Apply window/level\n print(\"Applying window/level adjustments...\")\n pixel_array = self.apply_window_level(pixel_array, ds)\n print(f\"Adjusted data range: {pixel_array.min()} to {pixel_array.max()}\")\n \n print(f\"Image shape: {pixel_array.shape}\")\n print(f\"Original dtype: {original_dtype}\")\n \n # Save metadata\n progress.setLabelText(\"Saving metadata...\")\n progress.setValue(60)\n \n metadata_file = os.path.join(self.output_directory, \n os.path.splitext(os.path.basename(self.input_file))[0] + \n \"_metadata.json\")\n with open(metadata_file, 'w') as f:\n json.dump(series_metadata, f, indent=2)\n \n # Get physical sizes from metadata\n pixel_spacing = series_metadata.get(\"PixelSpacing\", [1, 1])\n slice_thickness = series_metadata.get(\"SliceThickness\", 1)\n \n print(f\"Pixel spacing: {pixel_spacing}\")\n print(f\"Slice thickness: {slice_thickness}\")\n \n # Save TIFF\n progress.setLabelText(\"Saving TIFF file(s)...\")\n progress.setValue(80)\n \n # Convert back to original dtype if needed\n if np.issubdtype(original_dtype, np.integer):\n print(\"Converting back to original integer dtype...\")\n data_min = pixel_array.min()\n data_max = pixel_array.max()\n \n if data_max != data_min:\n pixel_array = ((pixel_array - data_min) / (data_max - data_min) * \n np.iinfo(original_dtype).max).astype(original_dtype)\n else:\n pixel_array = np.zeros_like(pixel_array, dtype=original_dtype)\n \n print(f\"Final data range: {pixel_array.min()} to {pixel_array.max()}\")\n \n # Prepare ImageJ metadata\n imagej_metadata = {\n 'axes': 'YX', # Will be updated to ZYX for 3D data\n 'spacing': float(slice_thickness), # Only used for 3D data\n 'unit': 'um',\n 'finterval': float(pixel_spacing[0]) # XY pixel size\n }\n \n base_name = os.path.splitext(os.path.basename(self.input_file))[0]\n \n if self.stack_radio.isChecked():\n # Save as single stack\n output_file = os.path.join(self.output_directory, f\"{base_name}.tif\")\n \n # Update axes for 3D data\n if len(pixel_array.shape) > 2:\n imagej_metadata['axes'] = 'ZYX'\n \n print(f\"Saving stack with metadata: {imagej_metadata}\")\n \n tifffile.imwrite(\n output_file,\n pixel_array,\n imagej=True,\n metadata=imagej_metadata,\n resolution=(1.0/float(pixel_spacing[0]), 1.0/float(pixel_spacing[1]))\n )\n \n print(f\"Saved stack to: {output_file}\")\n print(f\"Stack shape: {pixel_array.shape}\")\n \n # Replace the individual slices saving section in convert_dicom method with this:\n else:\n # For multi-slice DICOM, save individual slices\n if len(pixel_array.shape) > 2:\n imagej_metadata['axes'] = 'YX' # Reset to 2D for individual slices\n \n total_slices = pixel_array.shape[0]\n for i in range(total_slices):\n progress.setLabelText(f\"Saving slice {i+1}/{total_slices}...\")\n # Fix: Convert float to integer for progress value\n progress_value = int(80 + (i/total_slices)*15)\n progress.setValue(progress_value)\n QApplication.processEvents()\n \n if progress.wasCanceled():\n print(\"Operation cancelled by user\")\n return\n \n output_file = os.path.join(self.output_directory, \n f\"{base_name}_slice_{i+1:03d}.tif\")\n \n print(f\"Saving slice {i+1} with metadata: {imagej_metadata}\")\n \n tifffile.imwrite(\n output_file,\n pixel_array[i],\n imagej=True,\n metadata=imagej_metadata,\n resolution=(1.0/float(pixel_spacing[0]), 1.0/float(pixel_spacing[1]))\n )\n \n print(f\"Saved {total_slices} individual slices\")\n \n else:\n # Single slice DICOM\n output_file = os.path.join(self.output_directory, f\"{base_name}.tif\")\n \n print(f\"Saving single slice with metadata: {imagej_metadata}\")\n \n tifffile.imwrite(\n output_file,\n pixel_array,\n imagej=True,\n metadata=imagej_metadata,\n resolution=(1.0/float(pixel_spacing[0]), 1.0/float(pixel_spacing[1]))\n )\n \n print(f\"Saved single slice to: {output_file}\")\n \n progress.setValue(100)\n \n # Construct success message\n msg = \"Conversion complete!\\n\\n\"\n msg += f\"DICOM file: {os.path.basename(self.input_file)}\\n\"\n msg += f\"Output directory: {self.truncate_path(self.output_directory)}\\n\\n\"\n \n if self.stack_radio.isChecked():\n msg += f\"Saved as: {os.path.basename(output_file)}\\n\"\n else:\n if len(pixel_array.shape) > 2:\n msg += f\"Saved {pixel_array.shape[0]} individual slices\\n\"\n else:\n msg += f\"Saved as: {os.path.basename(output_file)}\\n\"\n \n msg += f\"\\nMetadata saved as: {os.path.basename(metadata_file)}\\n\"\n msg += f\"Pixel spacing: {pixel_spacing[0]}x{pixel_spacing[1]} µm\\n\"\n if len(pixel_array.shape) > 2:\n msg += f\"Slice thickness: {slice_thickness} µm\"\n \n QMessageBox.information(self, \"Success\", msg)\n \n except Exception as e:\n QMessageBox.critical(self, \"Error\", str(e))\n print(f\"Error occurred: {str(e)}\")\n import traceback\n traceback.print_exc()\n\n\n def show_centered(self, parent):\n parent_geo = parent.geometry()\n self.move(parent_geo.center() - self.rect().center())\n self.show()\n QApplication.processEvents() # Ensure window displays properly\n \ndef show_dicom_converter(parent):\n dialog = DicomConverter(parent)\n dialog.show_centered(parent)\n return dialog\n"
},
{
"path": "src/digitalsreeni_image_annotator/export_formats.py",
"content": "import json\nfrom PyQt5.QtGui import QImage\nfrom .utils import calculate_area, calculate_bbox\nimport yaml\nimport os\nimport shutil\nimport tempfile\nimport xml.etree.ElementTree as ET\nfrom xml.dom import minidom\nfrom datetime import datetime\n\nimport numpy as np\nimport skimage.draw\nfrom PIL import Image\n\n\n# Utility function to handle the COCO conversion for all export formats\ndef convert_to_coco(all_annotations, class_mapping, image_paths, slices, image_slices):\n with tempfile.TemporaryDirectory() as temp_dir:\n json_file_path, images_dir = export_coco_json(all_annotations, class_mapping, image_paths, slices, image_slices, temp_dir)\n \n with open(json_file_path, 'r') as f:\n coco_data = json.load(f)\n \n return coco_data, images_dir\n\n\n\ndef export_coco_json(all_annotations, class_mapping, image_paths, slices, image_slices, output_dir, json_filename=None):\n coco_format = {\n \"images\": [],\n \"categories\": [{\"id\": id, \"name\": name} for name, id in class_mapping.items()],\n \"annotations\": []\n }\n \n # Create images directory\n images_dir = os.path.join(output_dir, 'images')\n os.makedirs(images_dir, exist_ok=True)\n \n annotation_id = 1\n image_id = 1\n # Create a mapping of slice names to their QImage objects\n slice_map = {slice_name: qimage for slice_name, qimage in slices}\n \n # Handle all images and slices\n for image_name, annotations in all_annotations.items():\n # Skip if there are no annotations for this image/slice\n if not annotations:\n continue\n\n # Check if it's a slice (either in slice_map or has underscores and no file extension)\n is_slice = image_name in slice_map or ('_' in image_name and '.' not in image_name)\n \n if is_slice:\n qimage = slice_map.get(image_name)\n if qimage is None:\n # If the slice is not in slice_map, it might be a CZI slice or a TIFF slice\n # Find the corresponding QImage in slices or image_slices\n matching_slices = [s for s in slices if s[0] == image_name]\n if matching_slices:\n qimage = matching_slices[0][1]\n else:\n # Check in image_slices\n for stack_slices in image_slices.values():\n matching_slices = [s for s in stack_slices if s[0] == image_name]\n if matching_slices:\n qimage = matching_slices[0][1]\n break\n if qimage is None:\n print(f\"No image data found for slice {image_name}, skipping\")\n continue\n file_name_img = f\"{image_name}.png\"\n # Save the QImage as a file\n save_path = os.path.join(images_dir, file_name_img)\n if not os.path.exists(save_path):\n qimage.save(save_path)\n else:\n print(f\"Image {file_name_img} already exists in the target directory. Skipping save.\")\n else:\n # Check if the image_name exists in image_paths\n image_path = next((path for name, path in image_paths.items() if image_name in name), None)\n if not image_path:\n print(f\"No image path found for {image_name}, skipping\")\n continue\n if image_path.lower().endswith(('.tif', '.tiff', '.czi')):\n print(f\"Skipping main tiff/czi file: {image_name}\")\n continue\n file_name_img = image_name\n # Copy the image file\n dst_path = os.path.join(images_dir, file_name_img)\n if not os.path.exists(dst_path):\n shutil.copy2(image_path, dst_path)\n else:\n print(f\"Image {file_name_img} already exists in the target directory. Skipping copy.\")\n\n image_info = {\n \"file_name\": file_name_img,\n \"height\": qimage.height() if is_slice else QImage(image_path).height(),\n \"width\": qimage.width() if is_slice else QImage(image_path).width(),\n \"id\": image_id\n }\n coco_format[\"images\"].append(image_info)\n \n for class_name, class_annotations in annotations.items():\n for ann in class_annotations:\n coco_ann = create_coco_annotation(ann, image_id, annotation_id, class_name, class_mapping)\n coco_format[\"annotations\"].append(coco_ann)\n annotation_id += 1\n \n image_id += 1\n\n # Generate JSON filename if not provided\n if json_filename is None:\n timestamp = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n json_filename = f\"annotations_{timestamp}.json\"\n elif not json_filename.lower().endswith('.json'):\n json_filename += '.json'\n\n # Save COCO JSON file\n json_file_path = os.path.join(output_dir, json_filename)\n with open(json_file_path, 'w') as f:\n json.dump(coco_format, f, indent=2)\n\n return json_file_path, images_dir\n\n\ndef create_coco_annotation(ann, image_id, annotation_id, class_name, class_mapping):\n coco_ann = {\n \"id\": annotation_id,\n \"image_id\": image_id,\n \"category_id\": class_mapping[class_name],\n \"area\": calculate_area(ann),\n \"iscrowd\": 0\n }\n \n if \"segmentation\" in ann:\n coco_ann[\"segmentation\"] = [ann[\"segmentation\"]]\n coco_ann[\"bbox\"] = calculate_bbox(ann[\"segmentation\"])\n elif \"bbox\" in ann:\n coco_ann[\"bbox\"] = ann[\"bbox\"]\n \n return coco_ann\n\n\n\ndef export_yolo_v4(all_annotations, class_mapping, image_paths, slices, image_slices, output_dir):\n # Create output directories\n train_dir = os.path.join(output_dir, 'train')\n valid_dir = os.path.join(output_dir, 'valid')\n for dir_path in [train_dir, valid_dir]:\n os.makedirs(os.path.join(dir_path, 'images'), exist_ok=True)\n os.makedirs(os.path.join(dir_path, 'labels'), exist_ok=True)\n\n # Create a mapping of class names to YOLO indices\n class_to_index = {name: i for i, name in enumerate(class_mapping.keys())}\n\n # Create a mapping of slice names to their QImage objects\n slice_map = {slice_name: qimage for slice_name, qimage in slices}\n\n for image_name, annotations in all_annotations.items():\n # Skip if there are no annotations for this image/slice\n if not annotations:\n continue\n\n # For simplicity, we'll put all data in the train directory\n images_dir = os.path.join(train_dir, 'images')\n labels_dir = os.path.join(train_dir, 'labels')\n\n # Handle image saving (similar to before, but adjusted for new directory structure)\n if image_name in slice_map or ('_' in image_name and '.' not in image_name):\n # Handle slice images\n qimage = slice_map.get(image_name) or next((s[1] for s in slices if s[0] == image_name), None)\n if qimage is None:\n for stack_slices in image_slices.values():\n qimage = next((s[1] for s in stack_slices if s[0] == image_name), None)\n if qimage:\n break\n if qimage is None:\n print(f\"No image data found for slice {image_name}, skipping\")\n continue\n file_name_img = f\"{image_name}.png\"\n save_path = os.path.join(images_dir, file_name_img)\n if not os.path.exists(save_path):\n qimage.save(save_path)\n img_width, img_height = qimage.width(), qimage.height()\n else:\n # Handle regular images\n image_path = next((path for name, path in image_paths.items() if image_name in name), None)\n if not image_path or image_path.lower().endswith(('.tif', '.tiff', '.czi')):\n print(f\"Skipping file: {image_name}\")\n continue\n file_name_img = image_name\n dst_path = os.path.join(images_dir, file_name_img)\n if not os.path.exists(dst_path):\n shutil.copy2(image_path, dst_path)\n img = QImage(image_path)\n img_width, img_height = img.width(), img.height()\n\n # Write YOLO format annotation\n label_file = os.path.splitext(file_name_img)[0] + '.txt'\n with open(os.path.join(labels_dir, label_file), 'w') as f:\n for class_name, class_annotations in annotations.items():\n class_index = class_to_index[class_name]\n for ann in class_annotations:\n if 'segmentation' in ann:\n polygon = ann['segmentation']\n normalized_polygon = [coord / img_width if i % 2 == 0 else coord / img_height for i, coord in enumerate(polygon)]\n f.write(f\"{class_index} \" + \" \".join(map(lambda x: f\"{x:.6f}\", normalized_polygon)) + \"\\n\")\n elif 'bbox' in ann:\n x, y, w, h = ann['bbox']\n x_center = (x + w/2) / img_width\n y_center = (y + h/2) / img_height\n w = w / img_width\n h = h / img_height\n f.write(f\"{class_index} {x_center:.6f} {y_center:.6f} {w:.6f} {h:.6f}\\n\")\n\n # Create YAML file\n names = list(class_mapping.keys())\n yaml_data = {\n 'train': os.path.abspath(os.path.join(train_dir, 'images')),\n 'val': os.path.abspath(os.path.join(train_dir, 'images')), # Using train as val\n 'test': '../test/images', # Placeholder\n 'nc': len(names),\n 'names': names\n }\n\n # Save YAML file in the output directory\n yaml_path = os.path.join(output_dir, 'data.yaml')\n with open(yaml_path, 'w') as f:\n yaml.dump(yaml_data, f, default_flow_style=False)\n\n return train_dir, yaml_path\n\n\n\ndef export_yolo_v5plus(all_annotations, class_mapping, image_paths, slices, image_slices, output_dir):\n \"\"\"\n Export annotations in YOLO v5+ format.\n Directory structure:\n output_dir/\n ├── data.yaml\n ├── images/\n │ ├── train/\n │ └── val/\n └── labels/\n ├── train/\n └── val/\n \"\"\"\n # Create output directories with new structure\n images_train_dir = os.path.join(output_dir, 'images', 'train')\n images_val_dir = os.path.join(output_dir, 'images', 'val')\n labels_train_dir = os.path.join(output_dir, 'labels', 'train')\n labels_val_dir = os.path.join(output_dir, 'labels', 'val')\n\n for dir_path in [images_train_dir, images_val_dir, labels_train_dir, labels_val_dir]:\n os.makedirs(dir_path, exist_ok=True)\n\n # Create a mapping of class names to YOLO indices\n class_to_index = {name: i for i, name in enumerate(class_mapping.keys())}\n\n # Create a mapping of slice names to their QImage objects\n slice_map = {slice_name: qimage for slice_name, qimage in slices}\n\n for image_name, annotations in all_annotations.items():\n # Skip if there are no annotations for this image/slice\n if not annotations:\n continue\n\n # For simplicity, we'll put all data in the train directory\n # In practice, you might want to implement train/val split logic\n images_dir = images_train_dir\n labels_dir = labels_train_dir\n\n # Handle image saving (similar logic to the v4 version)\n if image_name in slice_map or ('_' in image_name and '.' not in image_name):\n # Handle slice images\n qimage = slice_map.get(image_name)\n if qimage is None:\n for stack_slices in image_slices.values():\n qimage = next((s[1] for s in stack_slices if s[0] == image_name), None)\n if qimage:\n break\n if qimage is None:\n print(f\"No image data found for slice {image_name}, skipping\")\n continue\n file_name_img = f\"{image_name}.png\"\n save_path = os.path.join(images_dir, file_name_img)\n if not os.path.exists(save_path):\n qimage.save(save_path)\n img_width, img_height = qimage.width(), qimage.height()\n else:\n # Handle regular images\n image_path = next((path for name, path in image_paths.items() if image_name in name), None)\n if not image_path or image_path.lower().endswith(('.tif', '.tiff', '.czi')):\n print(f\"Skipping file: {image_name}\")\n continue\n file_name_img = image_name\n dst_path = os.path.join(images_dir, file_name_img)\n if not os.path.exists(dst_path):\n shutil.copy2(image_path, dst_path)\n img = QImage(image_path)\n img_width, img_height = img.width(), img.height()\n\n # Write YOLO format annotation\n label_file = os.path.splitext(file_name_img)[0] + '.txt'\n with open(os.path.join(labels_dir, label_file), 'w') as f:\n for class_name, class_annotations in annotations.items():\n class_index = class_to_index[class_name]\n for ann in class_annotations:\n if 'segmentation' in ann:\n polygon = ann['segmentation']\n normalized_polygon = [coord / img_width if i % 2 == 0 else coord / img_height \n for i, coord in enumerate(polygon)]\n f.write(f\"{class_index} \" + \" \".join(map(lambda x: f\"{x:.6f}\", normalized_polygon)) + \"\\n\")\n elif 'bbox' in ann:\n x, y, w, h = ann['bbox']\n x_center = (x + w/2) / img_width\n y_center = (y + h/2) / img_height\n w = w / img_width\n h = h / img_height\n f.write(f\"{class_index} {x_center:.6f} {y_center:.6f} {w:.6f} {h:.6f}\\n\")\n\n # Create YAML file\n names = list(class_mapping.keys())\n yaml_data = {\n 'path': os.path.abspath(output_dir), # Root directory\n 'train': os.path.join('images', 'train'), # Relative to path\n 'val': os.path.join('images', 'val'), # Relative to path\n 'nc': len(names),\n 'names': names\n }\n\n # Save YAML file in the output directory\n yaml_path = os.path.join(output_dir, 'data.yaml')\n with open(yaml_path, 'w') as f:\n yaml.dump(yaml_data, f, default_flow_style=False)\n\n return output_dir, yaml_path\n\n\n\ndef export_labeled_images(all_annotations, class_mapping, image_paths, slices, image_slices, output_dir):\n # Create output directories\n images_dir = os.path.join(output_dir, 'images')\n labeled_images_dir = os.path.join(output_dir, 'labeled_images')\n os.makedirs(images_dir, exist_ok=True)\n os.makedirs(labeled_images_dir, exist_ok=True)\n\n # Create a dictionary to store class information for the summary\n class_summary = {class_name: [] for class_name in class_mapping.keys()}\n\n # Create directories for each class inside labeled_images_dir\n for class_name in class_mapping.keys():\n os.makedirs(os.path.join(labeled_images_dir, class_name), exist_ok=True)\n\n # Create a mapping of slice names to their QImage objects\n slice_map = {slice_name: qimage for slice_name, qimage in slices}\n\n for image_name, annotations in all_annotations.items():\n # Skip if there are no annotations for this image/slice\n if not annotations:\n continue\n\n # Check if it's a slice (either in slice_map or has underscores and no file extension)\n is_slice = image_name in slice_map or ('_' in image_name and '.' not in image_name)\n \n if is_slice:\n qimage = slice_map.get(image_name)\n if qimage is None:\n # If the slice is not in slice_map, it might be a CZI slice or a TIFF slice\n matching_slices = [s for s in slices if s[0] == image_name]\n if matching_slices:\n qimage = matching_slices[0][1]\n else:\n # Check in image_slices\n for stack_slices in image_slices.values():\n matching_slices = [s for s in stack_slices if s[0] == image_name]\n if matching_slices:\n qimage = matching_slices[0][1]\n break\n if qimage is None:\n print(f\"No image data found for slice {image_name}, skipping\")\n continue\n file_name_img = f\"{image_name}.png\"\n # Save the QImage as a file\n save_path = os.path.join(images_dir, file_name_img)\n if not os.path.exists(save_path):\n qimage.save(save_path)\n else:\n print(f\"Image {file_name_img} already exists in the target directory. Skipping copy.\")\n img_width, img_height = qimage.width(), qimage.height()\n else:\n # Check if the image_name exists in image_paths\n image_path = next((path for name, path in image_paths.items() if image_name in name), None)\n if not image_path:\n print(f\"No image path found for {image_name}, skipping\")\n continue\n if image_path.lower().endswith(('.tif', '.tiff', '.czi')):\n print(f\"Skipping main tiff/czi file: {image_name}\")\n continue\n file_name_img = image_name\n # Copy the image file\n dst_path = os.path.join(images_dir, file_name_img)\n if not os.path.exists(dst_path):\n shutil.copy2(image_path, dst_path)\n else:\n print(f\"Image {file_name_img} already exists in the target directory. Skipping copy.\")\n\n\n img = Image.open(image_path)\n img_width, img_height = img.size\n\n # Create a dictionary to store masks for each class\n class_masks = {class_name: np.zeros((img_height, img_width), dtype=np.uint16) for class_name in class_mapping.keys()}\n\n for class_name, class_annotations in annotations.items():\n mask = class_masks[class_name]\n for ann in class_annotations:\n object_number = np.max(mask) + 1 # Increment object number for this class\n \n if 'segmentation' in ann:\n polygon = np.array(ann['segmentation']).reshape(-1, 2)\n rr, cc = skimage.draw.polygon(polygon[:, 1], polygon[:, 0], (img_height, img_width))\n mask[rr, cc] = object_number\n elif 'bbox' in ann:\n x, y, w, h = map(int, ann['bbox'])\n mask[y:y+h, x:x+w] = object_number\n\n class_summary[class_name].append(file_name_img)\n\n # Save masks for each class\n for class_name, mask in class_masks.items():\n if np.any(mask): # Only save if the mask is not empty\n mask_filename = f\"{os.path.splitext(file_name_img)[0]}_{class_name}_mask.png\"\n mask_path = os.path.join(labeled_images_dir, class_name, mask_filename)\n Image.fromarray(mask.astype(np.uint16)).save(mask_path)\n\n # Create summary text file\n summary_path = os.path.join(labeled_images_dir, 'class_summary.txt')\n with open(summary_path, 'w') as f:\n f.write(\"Classes (folder names):\\n\")\n for class_name, files in class_summary.items():\n if files: # Only include classes that have annotations\n f.write(f\"- {class_name}\\n\")\n f.write(f\" Images: {', '.join(sorted(set(files)))}\\n\\n\")\n\n return output_dir\n\n\n\ndef export_semantic_labels(all_annotations, class_mapping, image_paths, slices, image_slices, output_dir):\n # Create output directories\n images_dir = os.path.join(output_dir, 'images')\n segmented_images_dir = os.path.join(output_dir, 'segmented_images')\n os.makedirs(images_dir, exist_ok=True)\n os.makedirs(segmented_images_dir, exist_ok=True)\n\n # Create a mapping of class names to unique pixel values\n class_to_pixel = {name: i+1 for i, name in enumerate(sorted(class_mapping.keys()))}\n\n # Create a mapping of slice names to their QImage objects\n slice_map = {slice_name: qimage for slice_name, qimage in slices}\n\n for image_name, annotations in all_annotations.items():\n # Skip if there are no annotations for this image/slice\n if not annotations:\n continue\n\n # Check if it's a slice (either in slice_map or has underscores and no file extension)\n is_slice = image_name in slice_map or ('_' in image_name and '.' not in image_name)\n \n if is_slice:\n qimage = slice_map.get(image_name)\n if qimage is None:\n # If the slice is not in slice_map, it might be a CZI slice or a TIFF slice\n matching_slices = [s for s in slices if s[0] == image_name]\n if matching_slices:\n qimage = matching_slices[0][1]\n else:\n # Check in image_slices\n for stack_slices in image_slices.values():\n matching_slices = [s for s in stack_slices if s[0] == image_name]\n if matching_slices:\n qimage = matching_slices[0][1]\n break\n if qimage is None:\n print(f\"No image data found for slice {image_name}, skipping\")\n continue\n file_name_img = f\"{image_name}.png\"\n # Save the QImage as a file\n save_path = os.path.join(images_dir, file_name_img)\n if not os.path.exists(save_path):\n qimage.save(save_path)\n else:\n print(f\"Image {file_name_img} already exists in the target directory. Skipping copy.\")\n img_width, img_height = qimage.width(), qimage.height()\n else:\n # Check if the image_name exists in image_paths\n image_path = next((path for name, path in image_paths.items() if image_name in name), None)\n if not image_path:\n print(f\"No image path found for {image_name}, skipping\")\n continue\n if image_path.lower().endswith(('.tif', '.tiff', '.czi')):\n print(f\"Skipping main tiff/czi file: {image_name}\")\n continue\n file_name_img = image_name\n # Copy the image file\n dst_path = os.path.join(images_dir, file_name_img)\n if not os.path.exists(dst_path):\n shutil.copy2(image_path, dst_path)\n else:\n print(f\"Image {file_name_img} already exists in the target directory. Skipping copy.\")\n\n img = Image.open(image_path)\n img_width, img_height = img.size\n\n # Create a single mask for all classes\n semantic_mask = np.zeros((img_height, img_width), dtype=np.uint8)\n\n for class_name, class_annotations in annotations.items():\n pixel_value = class_to_pixel[class_name]\n for ann in class_annotations:\n if 'segmentation' in ann:\n polygon = np.array(ann['segmentation']).reshape(-1, 2)\n rr, cc = skimage.draw.polygon(polygon[:, 1], polygon[:, 0], (img_height, img_width))\n semantic_mask[rr, cc] = pixel_value\n elif 'bbox' in ann:\n x, y, w, h = map(int, ann['bbox'])\n semantic_mask[y:y+h, x:x+w] = pixel_value\n\n # Save semantic mask\n mask_filename = f\"{os.path.splitext(file_name_img)[0]}_semantic_mask.png\"\n mask_path = os.path.join(segmented_images_dir, mask_filename)\n Image.fromarray(semantic_mask).save(mask_path)\n\n # Create class mapping text file\n mapping_path = os.path.join(segmented_images_dir, 'class_pixel_mapping.txt')\n with open(mapping_path, 'w') as f:\n f.write(\"Pixel Value : Class Name\\n\")\n for class_name, pixel_value in class_to_pixel.items():\n f.write(f\"{pixel_value} : {class_name}\\n\")\n\n return output_dir\n\n\n\ndef export_pascal_voc_bbox(all_annotations, class_mapping, image_paths, slices, image_slices, output_dir):\n # Create output directories\n images_dir = os.path.join(output_dir, 'images')\n annotations_dir = os.path.join(output_dir, 'Annotations')\n os.makedirs(images_dir, exist_ok=True)\n os.makedirs(annotations_dir, exist_ok=True)\n\n # Create a mapping of slice names to their QImage objects\n slice_map = {slice_name: qimage for slice_name, qimage in slices}\n\n for image_name, annotations in all_annotations.items():\n # Skip if there are no annotations for this image/slice\n if not annotations:\n continue\n\n # Check if it's a slice (either in slice_map or has underscores and no file extension)\n is_slice = image_name in slice_map or ('_' in image_name and '.' not in image_name)\n \n if is_slice:\n qimage = slice_map.get(image_name)\n if qimage is None:\n # If the slice is not in slice_map, it might be a CZI slice or a TIFF slice\n matching_slices = [s for s in slices if s[0] == image_name]\n if matching_slices:\n qimage = matching_slices[0][1]\n else:\n # Check in image_slices\n for stack_slices in image_slices.values():\n matching_slices = [s for s in stack_slices if s[0] == image_name]\n if matching_slices:\n qimage = matching_slices[0][1]\n break\n if qimage is None:\n print(f\"No image data found for slice {image_name}, skipping\")\n continue\n file_name_img = f\"{image_name}.png\"\n # Save the QImage as a file\n save_path = os.path.join(images_dir, file_name_img)\n if not os.path.exists(save_path):\n qimage.save(save_path)\n else:\n print(f\"Image {file_name_img} already exists in the target directory. Skipping copy.\")\n img_width, img_height = qimage.width(), qimage.height()\n else:\n # Check if the image_name exists in image_paths\n image_path = next((path for name, path in image_paths.items() if image_name in name), None)\n if not image_path:\n print(f\"No image path found for {image_name}, skipping\")\n continue\n if image_path.lower().endswith(('.tif', '.tiff', '.czi')):\n print(f\"Skipping main tiff/czi file: {image_name}\")\n continue\n file_name_img = image_name\n # Copy the image file\n dst_path = os.path.join(images_dir, file_name_img)\n if not os.path.exists(dst_path):\n shutil.copy2(image_path, dst_path)\n else:\n print(f\"Image {file_name_img} already exists in the target directory. Skipping copy.\")\n\n img = QImage(image_path)\n img_width, img_height = img.width(), img.height()\n\n # Create the XML structure\n root = ET.Element('annotation')\n ET.SubElement(root, 'folder').text = 'images'\n ET.SubElement(root, 'filename').text = file_name_img\n ET.SubElement(root, 'path').text = os.path.join('images', file_name_img)\n\n size = ET.SubElement(root, 'size')\n ET.SubElement(size, 'width').text = str(img_width)\n ET.SubElement(size, 'height').text = str(img_height)\n ET.SubElement(size, 'depth').text = '3' # Assuming RGB images\n\n ET.SubElement(root, 'segmented').text = '0'\n\n # Add object annotations\n for class_name, class_annotations in annotations.items():\n for ann in class_annotations:\n obj = ET.SubElement(root, 'object')\n ET.SubElement(obj, 'name').text = class_name\n ET.SubElement(obj, 'pose').text = 'Unspecified'\n ET.SubElement(obj, 'truncated').text = '0'\n ET.SubElement(obj, 'difficult').text = '0'\n\n if 'bbox' in ann:\n x, y, w, h = ann['bbox']\n bndbox = ET.SubElement(obj, 'bndbox')\n ET.SubElement(bndbox, 'xmin').text = str(int(x))\n ET.SubElement(bndbox, 'ymin').text = str(int(y))\n ET.SubElement(bndbox, 'xmax').text = str(int(x + w))\n ET.SubElement(bndbox, 'ymax').text = str(int(y + h))\n \n # Save the XML file\n xml_str = minidom.parseString(ET.tostring(root)).toprettyxml(indent=\" \")\n xml_filename = os.path.splitext(file_name_img)[0] + '.xml'\n with open(os.path.join(annotations_dir, xml_filename), 'w') as f:\n f.write(xml_str)\n \n return output_dir \n\n\n\ndef export_pascal_voc_both(all_annotations, class_mapping, image_paths, slices, image_slices, output_dir):\n # Create output directories\n images_dir = os.path.join(output_dir, 'images')\n annotations_dir = os.path.join(output_dir, 'Annotations')\n os.makedirs(images_dir, exist_ok=True)\n os.makedirs(annotations_dir, exist_ok=True)\n\n # Create a mapping of slice names to their QImage objects\n slice_map = {slice_name: qimage for slice_name, qimage in slices}\n\n for image_name, annotations in all_annotations.items():\n # Skip if there are no annotations for this image/slice\n if not annotations:\n continue\n\n # Check if it's a slice (either in slice_map or has underscores and no file extension)\n is_slice = image_name in slice_map or ('_' in image_name and '.' not in image_name)\n \n if is_slice:\n qimage = slice_map.get(image_name)\n if qimage is None:\n # If the slice is not in slice_map, it might be a CZI slice or a TIFF slice\n matching_slices = [s for s in slices if s[0] == image_name]\n if matching_slices:\n qimage = matching_slices[0][1]\n else:\n # Check in image_slices\n for stack_slices in image_slices.values():\n matching_slices = [s for s in stack_slices if s[0] == image_name]\n if matching_slices:\n qimage = matching_slices[0][1]\n break\n if qimage is None:\n print(f\"No image data found for slice {image_name}, skipping\")\n continue\n file_name_img = f\"{image_name}.png\"\n # Save the QImage as a file\n save_path = os.path.join(images_dir, file_name_img)\n if not os.path.exists(save_path):\n qimage.save(save_path)\n else:\n print(f\"Image {file_name_img} already exists in the target directory. Skipping copy.\")\n img_width, img_height = qimage.width(), qimage.height()\n else:\n # Check if the image_name exists in image_paths\n image_path = next((path for name, path in image_paths.items() if image_name in name), None)\n if not image_path:\n print(f\"No image path found for {image_name}, skipping\")\n continue\n if image_path.lower().endswith(('.tif', '.tiff', '.czi')):\n print(f\"Skipping main tiff/czi file: {image_name}\")\n continue\n file_name_img = image_name\n # Copy the image file\n dst_path = os.path.join(images_dir, file_name_img)\n if not os.path.exists(dst_path):\n shutil.copy2(image_path, dst_path)\n else:\n print(f\"Image {file_name_img} already exists in the target directory. Skipping copy.\")\n\n img = QImage(image_path)\n img_width, img_height = img.width(), img.height()\n\n # Create the XML structure\n root = ET.Element('annotation')\n ET.SubElement(root, 'folder').text = 'images'\n ET.SubElement(root, 'filename').text = file_name_img\n ET.SubElement(root, 'path').text = os.path.join('images', file_name_img)\n\n size = ET.SubElement(root, 'size')\n ET.SubElement(size, 'width').text = str(img_width)\n ET.SubElement(size, 'height').text = str(img_height)\n ET.SubElement(size, 'depth').text = '3' # Assuming RGB images\n\n ET.SubElement(root, 'segmented').text = '1' # Set to 1 if segmentation is included\n\n # Add object annotations\n for class_name, class_annotations in annotations.items():\n for ann in class_annotations:\n obj = ET.SubElement(root, 'object')\n ET.SubElement(obj, 'name').text = class_name\n ET.SubElement(obj, 'pose').text = 'Unspecified'\n ET.SubElement(obj, 'truncated').text = '0'\n ET.SubElement(obj, 'difficult').text = '0'\n\n if 'bbox' in ann:\n x, y, w, h = ann['bbox']\n bndbox = ET.SubElement(obj, 'bndbox')\n ET.SubElement(bndbox, 'xmin').text = str(int(x))\n ET.SubElement(bndbox, 'ymin').text = str(int(y))\n ET.SubElement(bndbox, 'xmax').text = str(int(x + w))\n ET.SubElement(bndbox, 'ymax').text = str(int(y + h))\n\n if 'segmentation' in ann:\n segmentation = ET.SubElement(obj, 'segmentation')\n ET.SubElement(segmentation, 'area').text = str(ann.get('area', 0))\n \n # Convert polygon to a list of (x,y) tuples\n polygon = ann['segmentation']\n points = [(polygon[i], polygon[i+1]) for i in range(0, len(polygon), 2)]\n \n # Create the polygon element\n polygon_elem = ET.SubElement(segmentation, 'polygon')\n for i, (x, y) in enumerate(points):\n point = ET.SubElement(polygon_elem, f'pt{i+1}')\n ET.SubElement(point, 'x').text = str(int(x))\n ET.SubElement(point, 'y').text = str(int(y))\n\n # Save the XML file\n xml_str = minidom.parseString(ET.tostring(root)).toprettyxml(indent=\" \")\n xml_filename = os.path.splitext(file_name_img)[0] + '.xml'\n with open(os.path.join(annotations_dir, xml_filename), 'w') as f:\n f.write(xml_str)\n\n return output_dir"
},
{
"path": "src/digitalsreeni_image_annotator/help_window.py",
"content": "from PyQt5.QtWidgets import QDialog, QVBoxLayout, QTextBrowser\nfrom PyQt5.QtCore import Qt\nfrom .soft_dark_stylesheet import soft_dark_stylesheet\nfrom .default_stylesheet import default_stylesheet\n\nclass HelpWindow(QDialog):\n def __init__(self, dark_mode=False, font_size=10):\n super().__init__()\n self.setWindowTitle(\"Help\")\n self.setModal(False) # Make it non-modal\n self.setGeometry(100, 100, 800, 600)\n layout = QVBoxLayout()\n self.text_browser = QTextBrowser()\n self.text_browser.setOpenExternalLinks(True)\n layout.addWidget(self.text_browser)\n self.setLayout(layout)\n \n if dark_mode:\n self.setStyleSheet(soft_dark_stylesheet)\n else:\n self.setStyleSheet(default_stylesheet)\n \n self.font_size = font_size\n self.apply_font_size()\n self.load_help_content()\n \n def show_centered(self, parent):\n parent_geo = parent.geometry()\n self.move(parent_geo.center() - self.rect().center())\n self.show()\n \n def apply_font_size(self):\n self.setStyleSheet(f\"QWidget {{ font-size: {self.font_size}pt; }}\")\n font = self.text_browser.font()\n font.setPointSize(self.font_size)\n self.text_browser.setFont(font)\n\n def load_help_content(self):\n help_text = \"\"\"\n Image Annotator Help Guide
\n\n Overview
\n Image Annotator is a user-friendly GUI tool designed for generating masks for image segmentation and object detection. It allows users to create, edit, and save annotations in various formats, including COCO-style JSON, YOLO v8, and Pascal VOC. Annotations can be defined using manual tools like the polygon tool or in a semi-automated way with the assistance of the Segment Anything Model (SAM-2) pre-trained model. The tool supports multi-dimensional images such as TIFF stacks and CZI files and provides dark mode and adjustable application font sizes for enhanced GUI experience.
\n\n Key Features
\n \n - Semi-automated annotations with SAM-2 assistance (Segment Anything Model)
\n - Manual annotations with polygons and rectangles
\n - Save and load projects for continued work
\n - Import existing COCO JSON annotations with images
\n - Export annotations to various formats (COCO JSON, YOLO v8, Labeled images, Semantic labels, Pascal VOC)
\n - Handle multi-dimensional images (TIFF stacks and CZI files)
\n - Zoom and pan for detailed annotations
\n - Support for multiple classes with customizable colors
\n - User-friendly interface with intuitive controls
\n - Adjustable application font size
\n - Dark mode for comfortable viewing
\n - Support for common image formats (PNG, JPG, BMP) and multi-dimensional formats (TIFF, CZI)
\n - Pick appropriate pre-trained SAM2 model for flexible and improved semi-automated annotations
\n - Additional tools for dataset management and image processing
\n
\n\n Getting Started
\n Starting a New Project
\n \n - Click \"New Project\" or use Ctrl+N to start a new project.
\n - Click \"Add New Images\" to import multiple images you want to annotate, including TIFF stacks and CZI files.
\n - For multi-dimensional images, you'll be prompted to assign dimensions (e.g., T for time, Z for depth, C for channels).
\n - Use \"Add Classes\" to define classes of interest.
\n - Start annotating by selecting a class and using the Polygon, Rectangle Tool, or SAM-Assisted tool.
\n
\n\n Opening an Existing Project
\n \n - Click \"Open Project\" or use Ctrl+O to load a previously saved project.
\n - If there are any missing images, you'll be prompted to locate them on your drive. Located images will be automatically copied to the project directory.
\n - If you choose not to locate missing images, the annotations for those images will be removed.
\n
\n\n Importing Existing Annotations
\n \n - Click \"Import Annotations with Images\" to load existing COCO JSON annotations along with their corresponding images.
\n - Select the COCO JSON file. The images should be in the same directory as the JSON file.
\n - The annotations and images will be loaded into your current project.
\n
\n\n Annotation Process
\n \n - Select a Class: Choose the class you want to annotate from the class list.
\n - Choose a Tool: Select either the Polygon Tool, Rectangle Tool, or SAM-Assisted tool.
\n - Create Annotation:\n
\n - For Polygon Tool: Click around the object to define its boundary. Press Enter or click \"Finish Polygon\" when done.
\n - For Rectangle Tool: Click and drag to create a bounding box.
\n - For SAM-Assisted tool: \n
\n - Select a SAM model from the \"Pick a SAM Model\" dropdown. It's recommended to use smaller models like SAM2 tiny or SAM2 small for better performance.
\n - Note: When you select a model for the first time, the application needs to download it. This process may take a few seconds to a minute, depending on your internet connection speed. Subsequent uses of the same model will be faster as it will already be cached locally, in your working directory.
\n - Click the \"SAM-Assisted\" button to activate the tool.
\n - Draw a rectangle around objects of interest to allow SAM2 to automatically detect objects.
\n - SAM2 will provide various outputs with different scores, and only the top-scoring region will be displayed.
\n - If the desired result isn't achieved on the first try, draw again.
\n - For low-quality images where SAM2 may not auto-detect objects, manual tools may be necessary.
\n
\n \n
\n \n
\n\n Exporting Annotations
\n \n - Click \"Export Annotations\" to open the export dialog.
\n - Select the desired export format from the dropdown menu.
\n - Choose the export location and confirm to save the annotations in the selected format.
\n
\n\n Navigation and Viewing
\n \n - Zoom: Use the slider at the bottom of the image, or hold Ctrl and use the mouse wheel.
\n - Pan: Hold Ctrl, click the left mouse button, and move the mouse.
\n - Switch Images: Click on an image name in the image list on the right.
\n - Navigate Slices: For multi-dimensional images, use the slice list on the right to click through slices.
\n
\n\n Tools Menu
\n The Tools menu provides access to various useful tools for dataset management and image processing. Each tool opens an intuitive GUI to guide you through the process:
\n \n - Annotation Statistics: Provides statistical information about your annotations.
\n - COCO JSON Combiner: Allows you to combine multiple COCO JSON annotation files.
\n - Dataset Splitter: Helps you split your dataset into train, validation, and test sets.
\n - Stack to Slices: Converts multi-dimensional image stacks into individual 2D slices.
\n - Image Patcher: Splits large images into smaller patches with or without overlap.
\n - Image Augmenter: Applies various transformations to augment your image dataset.
\n - Slice Registration: Aligns images in a stack with advanced registration options.
\n - Stack Interpolator: Adjusts Z-spacing in image stacks. Excellent tool to generate isotropic volumes from non-isotropic data.
\n - DICOM Converter: Converts DICOM files to TIFF format.
\n
\n\n Keyboard Shortcuts
\n \n - Ctrl + N: Create a new project
\n - Ctrl + O: Open an existing project
\n - Ctrl + S: Save the current project
\n - Ctrl + W: Close the current project
\n - Ctrl + Shift + S: Open Annotation Statistics
\n - F1: Open this help window
\n - Ctrl + Wheel: Zoom in/out
\n - Esc: Cancel current annotation, exit edit mode, or exit SAM-assisted annotation
\n - Enter: Finish current annotation, exit edit mode, or accept SAM-generated mask
\n - Up/Down Arrow Keys: Navigate through slices in multi-dimensional images
\n
\n\n Getting Help
\n If you encounter any issues or have suggestions for improvement, please open an issue on our GitHub repository or contact the development team.
\n \"\"\"\n self.text_browser.setHtml(help_text)\n"
},
{
"path": "src/digitalsreeni_image_annotator/image_augmenter.py",
"content": "import os\nimport random\nimport cv2\nimport numpy as np\nimport json\n\nfrom PyQt5.QtWidgets import (QDialog, QVBoxLayout, QHBoxLayout, QPushButton, \n QFileDialog, QLabel, QMessageBox, QSpinBox, \n QCheckBox, QDoubleSpinBox, QProgressBar, QApplication)\nfrom PyQt5.QtCore import Qt\n\nclass ImageAugmenterDialog(QDialog):\n def __init__(self, parent=None):\n super().__init__(parent)\n self.setWindowTitle(\"Image Augmenter\")\n self.setGeometry(100, 100, 400, 600)\n self.setWindowFlags(self.windowFlags() | Qt.Window)\n self.setWindowModality(Qt.ApplicationModal)\n self.input_dir = \"\"\n self.output_dir = \"\"\n self.coco_file = \"\"\n self.coco_data = None\n self.initUI()\n\n def initUI(self):\n layout = QVBoxLayout()\n\n # Input directory selection\n input_layout = QHBoxLayout()\n self.input_label = QLabel(\"Input Directory: Not selected\")\n input_button = QPushButton(\"Select Input Directory\")\n input_button.clicked.connect(self.select_input_directory)\n input_layout.addWidget(self.input_label)\n input_layout.addWidget(input_button)\n layout.addLayout(input_layout)\n\n # Output directory selection\n output_layout = QHBoxLayout()\n self.output_label = QLabel(\"Output Directory: Not selected\")\n output_button = QPushButton(\"Select Output Directory\")\n output_button.clicked.connect(self.select_output_directory)\n output_layout.addWidget(self.output_label)\n output_layout.addWidget(output_button)\n layout.addLayout(output_layout)\n\n # Number of augmentations per image\n aug_count_layout = QHBoxLayout()\n aug_count_layout.addWidget(QLabel(\"Augmentations per image:\"))\n self.aug_count_spin = QSpinBox()\n self.aug_count_spin.setRange(1, 100)\n self.aug_count_spin.setValue(5)\n aug_count_layout.addWidget(self.aug_count_spin)\n layout.addLayout(aug_count_layout)\n \n \n # Add COCO JSON annotation augmentation checkbox and file selection\n self.coco_check = QCheckBox(\"Augment COCO JSON annotations\")\n self.coco_check.stateChanged.connect(self.toggle_elastic_deformation)\n layout.addWidget(self.coco_check)\n \n coco_layout = QHBoxLayout()\n self.coco_label = QLabel(\"COCO JSON File: Not selected\")\n coco_button = QPushButton(\"Select COCO JSON\")\n coco_button.clicked.connect(self.select_coco_json)\n coco_layout.addWidget(self.coco_label)\n coco_layout.addWidget(coco_button)\n layout.addLayout(coco_layout)\n\n # Transformations\n layout.addWidget(QLabel(\"Transformations:\"))\n\n self.rotate_check = QCheckBox(\"Rotate\")\n self.rotate_spin = QSpinBox()\n self.rotate_spin.setRange(-180, 180)\n self.rotate_spin.setValue(30)\n rotate_layout = QHBoxLayout()\n rotate_layout.addWidget(self.rotate_check)\n rotate_layout.addWidget(QLabel(\"Max degrees:\"))\n rotate_layout.addWidget(self.rotate_spin)\n layout.addLayout(rotate_layout)\n\n self.zoom_check = QCheckBox(\"Zoom\")\n self.zoom_spin = QDoubleSpinBox()\n self.zoom_spin.setRange(0.1, 2.0)\n self.zoom_spin.setValue(0.2)\n self.zoom_spin.setSingleStep(0.1)\n zoom_layout = QHBoxLayout()\n zoom_layout.addWidget(self.zoom_check)\n zoom_layout.addWidget(QLabel(\"Scale factor:\"))\n zoom_layout.addWidget(self.zoom_spin)\n layout.addLayout(zoom_layout)\n\n self.blur_check = QCheckBox(\"Gaussian Blur\")\n layout.addWidget(self.blur_check)\n\n self.brightness_contrast_check = QCheckBox(\"Random Brightness and Contrast\")\n layout.addWidget(self.brightness_contrast_check)\n\n self.sharpen_check = QCheckBox(\"Sharpen\")\n layout.addWidget(self.sharpen_check)\n \n # Flip transformation\n flip_layout = QHBoxLayout()\n self.flip_check = QCheckBox(\"Flip\")\n flip_layout.addWidget(self.flip_check)\n self.flip_horizontal_check = QCheckBox(\"Horizontal\")\n self.flip_vertical_check = QCheckBox(\"Vertical\")\n flip_layout.addWidget(self.flip_horizontal_check)\n flip_layout.addWidget(self.flip_vertical_check)\n self.flip_horizontal_check.stateChanged.connect(self.update_flip_check)\n self.flip_vertical_check.stateChanged.connect(self.update_flip_check)\n layout.addLayout(flip_layout)\n\n\n # Elastic Deformation\n self.elastic_check = QCheckBox(\"Elastic Deformation\")\n layout.addWidget(self.elastic_check)\n elastic_layout = QHBoxLayout()\n elastic_layout.addWidget(self.elastic_check)\n elastic_layout.addWidget(QLabel(\"Alpha:\"))\n self.elastic_alpha_spin = QSpinBox()\n self.elastic_alpha_spin.setRange(1, 1000)\n self.elastic_alpha_spin.setValue(500)\n elastic_layout.addWidget(self.elastic_alpha_spin)\n elastic_layout.addWidget(QLabel(\"Sigma:\"))\n self.elastic_sigma_spin = QSpinBox()\n self.elastic_sigma_spin.setRange(1, 100)\n self.elastic_sigma_spin.setValue(20)\n elastic_layout.addWidget(self.elastic_sigma_spin)\n layout.addLayout(elastic_layout)\n\n # Grayscale Conversion\n self.grayscale_check = QCheckBox(\"Convert to Grayscale\")\n layout.addWidget(self.grayscale_check)\n\n # Histogram Equalization\n self.hist_equalize_check = QCheckBox(\"Histogram Equalization\")\n layout.addWidget(self.hist_equalize_check)\n\n\n # Augment button\n self.augment_button = QPushButton(\"Start Augmentation\")\n self.augment_button.clicked.connect(self.start_augmentation)\n layout.addWidget(self.augment_button)\n\n # Progress bar\n self.progress_bar = QProgressBar()\n layout.addWidget(self.progress_bar)\n\n self.setLayout(layout)\n\n def select_input_directory(self):\n self.input_dir = QFileDialog.getExistingDirectory(self, \"Select Input Directory\")\n if self.input_dir:\n self.input_label.setText(f\"Input Directory: {os.path.basename(self.input_dir)}\")\n\n def select_output_directory(self):\n self.output_dir = QFileDialog.getExistingDirectory(self, \"Select Output Directory\")\n if self.output_dir:\n self.output_label.setText(f\"Output Directory: {os.path.basename(self.output_dir)}\")\n\n def update_flip_check(self, state):\n if self.flip_horizontal_check.isChecked() or self.flip_vertical_check.isChecked():\n self.flip_check.setChecked(True)\n else:\n self.flip_check.setChecked(False)\n \n def select_coco_json(self):\n self.coco_file, _ = QFileDialog.getOpenFileName(self, \"Select COCO JSON File\", \"\", \"JSON Files (*.json)\")\n if self.coco_file:\n self.coco_label.setText(f\"COCO JSON File: {os.path.basename(self.coco_file)}\")\n with open(self.coco_file, 'r') as f:\n self.coco_data = json.load(f)\n self.coco_check.setChecked(True) # Automatically check the box when a file is loaded\n \n def toggle_elastic_deformation(self, state):\n if state == Qt.Checked:\n self.elastic_check.setChecked(False)\n self.elastic_check.setEnabled(False)\n else:\n self.elastic_check.setEnabled(True)\n \n\n \n def start_augmentation(self):\n if not self.input_dir or not self.output_dir:\n QMessageBox.warning(self, \"Missing Directory\", \"Please select both input and output directories.\")\n return\n \n if self.coco_check.isChecked() and not self.coco_file:\n QMessageBox.warning(self, \"Missing COCO JSON\", \"Please select a COCO JSON file for annotation augmentation.\")\n return\n \n # Create 'images' subdirectory in the output directory\n images_output_dir = os.path.join(self.output_dir, \"images\")\n os.makedirs(images_output_dir, exist_ok=True)\n \n image_files = [f for f in os.listdir(self.input_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tif', '.tiff'))]\n total_augmentations = len(image_files) * self.aug_count_spin.value()\n \n self.progress_bar.setMaximum(total_augmentations)\n self.progress_bar.setValue(0)\n \n augmented_coco_data = {\n \"images\": [],\n \"annotations\": [],\n \"categories\": self.coco_data[\"categories\"] if self.coco_data else []\n }\n \n next_image_id = 1\n next_annotation_id = 1\n \n for i, image_file in enumerate(image_files):\n input_path = os.path.join(self.input_dir, image_file)\n image = cv2.imread(input_path, cv2.IMREAD_UNCHANGED)\n \n if image is None:\n print(f\"Error loading image: {input_path}\")\n continue\n \n # Determine image type and bit depth\n is_color = len(image.shape) == 3 and image.shape[2] == 3\n bit_depth = image.dtype\n \n original_annotations = []\n if self.coco_check.isChecked():\n original_annotations = [ann for ann in self.coco_data[\"annotations\"] \n if any(img['file_name'] == image_file and img['id'] == ann['image_id'] \n for img in self.coco_data[\"images\"])]\n \n for j in range(self.aug_count_spin.value()):\n try:\n augmented, transform_params = self.apply_random_augmentation(image, include_annotations=self.coco_check.isChecked())\n \n # Ensure the augmented image has the same properties as the input\n if not is_color and len(augmented.shape) == 3:\n augmented = cv2.cvtColor(augmented, cv2.COLOR_BGR2GRAY)\n elif is_color and len(augmented.shape) == 2:\n augmented = cv2.cvtColor(augmented, cv2.COLOR_GRAY2BGR)\n \n augmented = augmented.astype(bit_depth)\n \n output_filename = f\"{os.path.splitext(image_file)[0]}_aug_{j+1}{os.path.splitext(image_file)[1]}\"\n output_path = os.path.join(images_output_dir, output_filename)\n cv2.imwrite(output_path, augmented)\n \n if self.coco_check.isChecked():\n augmented_coco_data[\"images\"].append({\n \"id\": next_image_id,\n \"file_name\": output_filename,\n \"height\": augmented.shape[0],\n \"width\": augmented.shape[1]\n })\n \n for ann in original_annotations:\n augmented_ann = self.augment_annotation(ann, transform_params, augmented.shape[:2])\n augmented_ann[\"id\"] = next_annotation_id\n augmented_ann[\"image_id\"] = next_image_id\n augmented_coco_data[\"annotations\"].append(augmented_ann)\n next_annotation_id += 1\n \n next_image_id += 1\n \n self.progress_bar.setValue(i * self.aug_count_spin.value() + j + 1)\n QApplication.processEvents()\n \n except Exception as e:\n print(f\"Error processing {image_file} (augmentation {j+1}): {str(e)}\")\n continue # Skip this augmentation and continue with the next\n \n if self.coco_check.isChecked():\n output_coco_path = os.path.join(self.output_dir, \"augmented_annotations.json\")\n with open(output_coco_path, 'w') as f:\n json.dump(augmented_coco_data, f, indent=2)\n \n QMessageBox.information(self, \"Augmentation Complete\", \"Image and annotation augmentation has been completed successfully.\")\n\n \n\n\n def apply_random_augmentation(self, image, include_annotations=False):\n augmentations = []\n \n if self.rotate_check.isChecked():\n augmentations.append(self.rotate_image)\n if self.zoom_check.isChecked():\n augmentations.append(self.zoom_image)\n if self.blur_check.isChecked():\n augmentations.append(self.blur_image)\n if self.brightness_contrast_check.isChecked():\n augmentations.append(self.adjust_brightness_contrast)\n if self.sharpen_check.isChecked():\n augmentations.append(self.sharpen_image)\n if self.flip_check.isChecked():\n augmentations.append(self.flip_image)\n if self.elastic_check.isChecked() and not include_annotations:\n augmentations.append(self.elastic_transform)\n if self.grayscale_check.isChecked():\n augmentations.append(self.convert_to_grayscale)\n if self.hist_equalize_check.isChecked():\n augmentations.append(self.apply_histogram_equalization)\n \n if not augmentations:\n return image, {}\n \n aug_func = random.choice(augmentations)\n return aug_func(image)\n\n def rotate_image(self, image):\n angle = random.uniform(-self.rotate_spin.value(), self.rotate_spin.value())\n h, w = image.shape[:2]\n center = (w / 2, h / 2)\n M = cv2.getRotationMatrix2D(center, -angle, 1.0) # Negative angle for clockwise rotation\n rotated = cv2.warpAffine(image, M, (w, h), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT)\n return rotated, {\"type\": \"rotate\", \"angle\": angle, \"center\": center, \"matrix\": M}\n \n def zoom_image(self, image):\n scale = random.uniform(1, 1 + self.zoom_spin.value())\n h, w = image.shape[:2]\n center = (w / 2, h / 2)\n M = cv2.getRotationMatrix2D(center, 0, scale)\n zoomed = cv2.warpAffine(image, M, (w, h), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT)\n return zoomed, {\"type\": \"zoom\", \"scale\": scale, \"center\": center, \"matrix\": M}\n \n \n \n\n def blur_image(self, image):\n kernel_size = random.choice([3, 5, 7])\n blurred = cv2.GaussianBlur(image, (kernel_size, kernel_size), 0)\n return blurred, {\"type\": \"blur\", \"kernel_size\": kernel_size}\n \n def adjust_brightness_contrast(self, image):\n alpha = random.uniform(0.5, 1.5) # Contrast control\n beta = random.uniform(-30, 30) # Brightness control\n adjusted = cv2.convertScaleAbs(image, alpha=alpha, beta=beta)\n return adjusted, {\"type\": \"brightness_contrast\", \"alpha\": alpha, \"beta\": beta}\n \n def sharpen_image(self, image):\n kernel = np.array([[-1,-1,-1], [-1,9,-1], [-1,-1,-1]])\n sharpened = cv2.filter2D(image, -1, kernel)\n return sharpened, {\"type\": \"sharpen\"}\n \n\n \n\n \n def flip_image(self, image):\n flip_options = []\n if self.flip_horizontal_check.isChecked():\n flip_options.append(1) # Horizontal flip\n if self.flip_vertical_check.isChecked():\n flip_options.append(0) # Vertical flip\n if self.flip_horizontal_check.isChecked() and self.flip_vertical_check.isChecked():\n flip_options.append(-1) # Both horizontal and vertical\n\n if not flip_options:\n return image, {\"type\": \"flip\", \"flip_code\": None}\n\n flip_code = random.choice(flip_options)\n flipped = cv2.flip(image, flip_code)\n return flipped, {\"type\": \"flip\", \"flip_code\": flip_code}\n\n def elastic_transform(self, image):\n alpha = self.elastic_alpha_spin.value()\n sigma = self.elastic_sigma_spin.value()\n shape = image.shape[:2]\n random_state = np.random.RandomState(None)\n dx = random_state.rand(*shape) * 2 - 1\n dy = random_state.rand(*shape) * 2 - 1\n dx = cv2.GaussianBlur(dx, (0, 0), sigma) * alpha\n dy = cv2.GaussianBlur(dy, (0, 0), sigma) * alpha\n x, y = np.meshgrid(np.arange(shape[1]), np.arange(shape[0]))\n distorted_x = x + dx\n distorted_y = y + dy\n \n transformed = cv2.remap(image, distorted_x.astype(np.float32), distorted_y.astype(np.float32), \n interpolation=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT_101)\n \n return transformed, {\"type\": \"elastic\", \"dx\": dx, \"dy\": dy, \"shape\": shape}\n \n def convert_to_grayscale(self, image):\n if len(image.shape) == 2:\n return image, {\"type\": \"grayscale\"} # Already grayscale\n gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)\n gray_3channel = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR) # Convert back to 3 channels\n return gray_3channel, {\"type\": \"grayscale\"}\n \n def apply_histogram_equalization(self, image):\n def equalize_8bit(img):\n return cv2.equalizeHist(img)\n \n def equalize_16bit(img):\n # Equalize 16-bit image\n hist, bins = np.histogram(img.flatten(), 65536, [0, 65536])\n cdf = hist.cumsum()\n cdf_normalized = cdf * 65535 / cdf[-1]\n equalized = np.interp(img.flatten(), bins[:-1], cdf_normalized).reshape(img.shape)\n return equalized.astype(np.uint16)\n \n if len(image.shape) == 2: # Grayscale image\n if image.dtype == np.uint8:\n equalized = equalize_8bit(image)\n elif image.dtype == np.uint16:\n equalized = equalize_16bit(image)\n else:\n raise ValueError(f\"Unsupported image dtype: {image.dtype}\")\n return equalized, {\"type\": \"histogram_equalization\", \"mode\": \"grayscale\"}\n else: # Color image\n # Convert to YUV color space\n yuv = cv2.cvtColor(image, cv2.COLOR_BGR2YUV)\n \n # Equalize the Y channel\n if image.dtype == np.uint8:\n yuv[:,:,0] = equalize_8bit(yuv[:,:,0])\n elif image.dtype == np.uint16:\n yuv[:,:,0] = equalize_16bit(yuv[:,:,0])\n else:\n raise ValueError(f\"Unsupported image dtype: {image.dtype}\")\n \n # Convert back to BGR color space\n equalized = cv2.cvtColor(yuv, cv2.COLOR_YUV2BGR)\n return equalized, {\"type\": \"histogram_equalization\", \"mode\": \"color\"}\n\n\n def augment_annotation(self, annotation, transform_params, image_shape):\n augmented_ann = annotation.copy()\n \n if transform_params[\"type\"] == \"rotate\":\n angle = transform_params[\"angle\"]\n center = transform_params[\"center\"]\n matrix = transform_params[\"matrix\"]\n augmented_ann[\"segmentation\"] = [self.rotate_polygon(annotation[\"segmentation\"][0], angle, center, matrix)]\n elif transform_params[\"type\"] == \"zoom\":\n scale = transform_params[\"scale\"]\n center = transform_params[\"center\"]\n matrix = transform_params[\"matrix\"]\n augmented_ann[\"segmentation\"] = [self.scale_polygon(annotation[\"segmentation\"][0], scale, center, matrix)]\n elif transform_params[\"type\"] == \"flip\":\n flip_code = transform_params[\"flip_code\"]\n if flip_code is not None:\n augmented_ann[\"segmentation\"] = [self.flip_polygon(annotation[\"segmentation\"][0], flip_code, image_shape)]\n elif transform_params[\"type\"] == \"elastic\":\n dx = transform_params[\"dx\"]\n dy = transform_params[\"dy\"]\n shape = transform_params[\"shape\"]\n augmented_ann[\"segmentation\"] = [self.elastic_transform_polygon(annotation[\"segmentation\"][0], dx, dy, shape)]\n \n # Recalculate bbox and area for all transformation types\n if \"segmentation\" in augmented_ann and augmented_ann[\"segmentation\"]:\n augmented_ann[\"bbox\"] = self.get_bbox_from_polygon(augmented_ann[\"segmentation\"][0])\n augmented_ann[\"area\"] = int(self.calculate_polygon_area(augmented_ann[\"segmentation\"][0]))\n \n return augmented_ann\n\n\n \n def calculate_polygon_area(self, polygon):\n points = np.array(polygon).reshape(-1, 2)\n return 0.5 * np.abs(np.dot(points[:, 0], np.roll(points[:, 1], 1)) - np.dot(points[:, 1], np.roll(points[:, 0], 1)))\n\n\n def rotate_polygon(self, polygon, angle, center, matrix):\n points = np.array(polygon).reshape(-1, 2)\n ones = np.ones(shape=(len(points), 1))\n points_ones = np.hstack([points, ones])\n transformed_points = matrix.dot(points_ones.T).T\n return np.round(transformed_points).astype(int).flatten().tolist()\n \n def scale_polygon(self, polygon, scale, center, matrix):\n points = np.array(polygon).reshape(-1, 2)\n ones = np.ones(shape=(len(points), 1))\n points_ones = np.hstack([points, ones])\n transformed_points = matrix.dot(points_ones.T).T\n return np.round(transformed_points).astype(int).flatten().tolist()\n \n def flip_polygon(self, polygon, flip_code, image_shape):\n points = np.array(polygon).reshape(-1, 2)\n if flip_code == 0: # Vertical flip\n points[:, 1] = image_shape[0] - points[:, 1]\n elif flip_code == 1: # Horizontal flip\n points[:, 0] = image_shape[1] - points[:, 0]\n elif flip_code == -1: # Both\n points[:, 0] = image_shape[1] - points[:, 0]\n points[:, 1] = image_shape[0] - points[:, 1]\n return np.round(points).astype(int).flatten().tolist()\n \n\n\n def get_bbox_from_polygon(self, polygon):\n points = np.array(polygon).reshape(-1, 2)\n x_min, y_min = np.min(points, axis=0)\n x_max, y_max = np.max(points, axis=0)\n return [int(x_min), int(y_min), int(x_max - x_min), int(y_max - y_min)]\n\n\n\n\n def show_centered(self, parent):\n parent_geo = parent.geometry()\n self.move(parent_geo.center() - self.rect().center())\n self.show()\n\ndef show_image_augmenter(parent):\n dialog = ImageAugmenterDialog(parent)\n dialog.show_centered(parent)\n return dialog"
},
{
"path": "src/digitalsreeni_image_annotator/image_label.py",
"content": "\"\"\"\nImageLabel module for the Image Annotator application.\n\nThis module contains the ImageLabel class, which is responsible for\ndisplaying the image and handling annotation interactions.\n\n@DigitalSreeni\nDr. Sreenivas Bhattiprolu\n\"\"\"\n\nfrom PyQt5.QtWidgets import QLabel, QApplication, QMessageBox\nfrom PyQt5.QtGui import (QPainter, QPen, QColor, QFont, QPolygonF, QBrush, QPolygon,\n QPixmap, QImage, QWheelEvent, QMouseEvent, QKeyEvent)\nfrom PyQt5.QtCore import Qt, QPoint, QPointF, QRectF, QSize\nfrom PIL import Image\nimport os\nimport warnings\nimport cv2\nimport numpy as np\n\nwarnings.filterwarnings(\"ignore\", category=UserWarning)\n\n\n\n\nclass ImageLabel(QLabel):\n \"\"\"\n A custom QLabel for displaying images and handling annotations.\n \"\"\"\n\n def __init__(self, parent=None):\n super().__init__(parent)\n self.annotations = {}\n self.current_annotation = []\n self.temp_point = None\n self.current_tool = None\n self.zoom_factor = 1.0\n self.class_colors = {}\n self.class_visibility = {}\n self.start_point = None\n self.end_point = None\n self.highlighted_annotations = []\n self.setMouseTracking(True)\n self.setFocusPolicy(Qt.StrongFocus)\n self.original_pixmap = None\n self.scaled_pixmap = None\n self.pan_start_pos = None\n self.main_window = None\n self.offset_x = 0\n self.offset_y = 0\n self.drawing_polygon = False\n self.editing_polygon = None\n self.editing_point_index = None\n self.hover_point_index = None\n self.fill_opacity = 0.3\n self.drawing_rectangle = False\n self.current_rectangle = None\n self.bit_depth = None\n self.image_path = None\n self.dark_mode = False\n \n self.paint_mask = None\n self.eraser_mask = None\n self.temp_paint_mask = None\n self.is_painting = False\n self.temp_eraser_mask = None\n self.is_erasing = False\n self.cursor_pos = None\n\n #SAM\n self.sam_magic_wand_active = False\n self.sam_bbox = None\n self.drawing_sam_bbox = False\n self.temp_sam_prediction = None\n \n self.temp_annotations = []\n\n\n\n \n def set_main_window(self, main_window):\n self.main_window = main_window\n \n def set_dark_mode(self, is_dark):\n self.dark_mode = is_dark\n self.update()\n\n def setPixmap(self, pixmap):\n \"\"\"Set the pixmap and update the scaled version.\"\"\"\n if isinstance(pixmap, QImage):\n pixmap = QPixmap.fromImage(pixmap)\n self.original_pixmap = pixmap\n self.update_scaled_pixmap()\n \n def detect_bit_depth(self):\n \"\"\"Detect and store the actual image bit depth using PIL.\"\"\"\n if self.image_path and os.path.exists(self.image_path):\n with Image.open(self.image_path) as img:\n if img.mode == '1':\n self.bit_depth = 1\n elif img.mode == 'L':\n self.bit_depth = 8\n elif img.mode == 'I;16':\n self.bit_depth = 16\n elif img.mode in ['RGB', 'HSV']:\n self.bit_depth = 24\n elif img.mode in ['RGBA', 'CMYK']:\n self.bit_depth = 32\n else:\n self.bit_depth = img.bits\n \n if self.main_window:\n self.main_window.update_image_info()\n\n def update_scaled_pixmap(self):\n if self.original_pixmap and not self.original_pixmap.isNull():\n scaled_size = self.original_pixmap.size() * self.zoom_factor\n self.scaled_pixmap = self.original_pixmap.scaled(\n scaled_size.width(),\n scaled_size.height(),\n Qt.KeepAspectRatio,\n Qt.SmoothTransformation\n )\n super().setPixmap(self.scaled_pixmap)\n self.setMinimumSize(self.scaled_pixmap.size())\n self.update_offset()\n else:\n self.scaled_pixmap = None\n super().setPixmap(QPixmap())\n self.setMinimumSize(QSize(0, 0))\n\n def update_offset(self):\n \"\"\"Update the offset for centered image display.\"\"\"\n if self.scaled_pixmap:\n self.offset_x = int((self.width() - self.scaled_pixmap.width()) / 2)\n self.offset_y = int((self.height() - self.scaled_pixmap.height()) / 2)\n \n def reset_annotation_state(self):\n \"\"\"Reset the annotation state.\"\"\"\n self.temp_point = None\n self.start_point = None\n self.end_point = None\n\n def clear_current_annotation(self):\n \"\"\"Clear the current annotation.\"\"\"\n self.current_annotation = []\n\n def resizeEvent(self, event):\n \"\"\"Handle resize events.\"\"\"\n super().resizeEvent(event)\n self.update_offset()\n \n \n def start_painting(self, pos):\n if self.temp_paint_mask is None:\n self.temp_paint_mask = np.zeros((self.original_pixmap.height(), self.original_pixmap.width()), dtype=np.uint8)\n self.is_painting = True\n self.continue_painting(pos)\n\n def continue_painting(self, pos):\n if not self.is_painting:\n return\n brush_size = self.main_window.paint_brush_size\n cv2.circle(self.temp_paint_mask, (int(pos[0]), int(pos[1])), brush_size, 255, -1)\n self.update()\n\n def finish_painting(self):\n if not self.is_painting:\n return\n self.is_painting = False\n # Don't commit the annotation yet, just keep the temp_paint_mask\n \n \n def commit_paint_annotation(self):\n if self.temp_paint_mask is not None and self.main_window.current_class:\n class_name = self.main_window.current_class\n contours, _ = cv2.findContours(self.temp_paint_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)\n for contour in contours:\n if cv2.contourArea(contour) > 10: # Minimum area threshold\n segmentation = contour.flatten().tolist()\n new_annotation = {\n \"segmentation\": segmentation,\n \"category_id\": self.main_window.class_mapping[class_name],\n \"category_name\": class_name,\n }\n self.annotations.setdefault(class_name, []).append(new_annotation)\n self.main_window.add_annotation_to_list(new_annotation)\n self.temp_paint_mask = None\n self.main_window.save_current_annotations()\n self.main_window.update_slice_list_colors()\n self.update()\n \n \n def discard_paint_annotation(self):\n self.temp_paint_mask = None\n self.update() \n \n\n def start_erasing(self, pos):\n if self.temp_eraser_mask is None:\n self.temp_eraser_mask = np.zeros((self.original_pixmap.height(), self.original_pixmap.width()), dtype=np.uint8)\n self.is_erasing = True\n self.continue_erasing(pos)\n\n def continue_erasing(self, pos):\n if not self.is_erasing:\n return\n eraser_size = self.main_window.eraser_size\n cv2.circle(self.temp_eraser_mask, (int(pos[0]), int(pos[1])), eraser_size, 255, -1)\n self.update()\n\n def finish_erasing(self):\n if not self.is_erasing:\n return\n self.is_erasing = False\n # Don't commit the eraser changes yet, just keep the temp_eraser_mask\n\n def commit_eraser_changes(self):\n if self.temp_eraser_mask is not None:\n eraser_mask = self.temp_eraser_mask.astype(bool)\n current_name = self.main_window.current_slice or self.main_window.image_file_name\n annotations_changed = False\n \n for class_name, annotations in self.annotations.items():\n updated_annotations = []\n max_number = max([ann.get('number', 0) for ann in annotations] + [0])\n for annotation in annotations:\n if \"segmentation\" in annotation and annotation[\"segmentation\"] is not None:\n points = np.array(annotation[\"segmentation\"]).reshape(-1, 2).astype(int)\n mask = np.zeros_like(self.temp_eraser_mask)\n cv2.fillPoly(mask, [points], 255)\n mask = mask.astype(bool)\n mask[eraser_mask] = False\n contours, _ = cv2.findContours(mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)\n for i, contour in enumerate(contours):\n if cv2.contourArea(contour) > 10: # Minimum area threshold\n new_segmentation = contour.flatten().tolist()\n new_annotation = annotation.copy()\n new_annotation[\"segmentation\"] = new_segmentation\n if i == 0:\n new_annotation[\"number\"] = annotation.get(\"number\", max_number + 1)\n else:\n max_number += 1\n new_annotation[\"number\"] = max_number\n updated_annotations.append(new_annotation)\n if len(contours) > 1:\n annotations_changed = True\n else:\n updated_annotations.append(annotation)\n self.annotations[class_name] = updated_annotations\n \n self.temp_eraser_mask = None\n \n # Update the all_annotations dictionary in the main window\n self.main_window.all_annotations[current_name] = self.annotations\n \n # Call update_annotation_list directly\n self.main_window.update_annotation_list()\n \n self.main_window.save_current_annotations()\n self.main_window.update_slice_list_colors()\n self.update()\n \n #print(f\"Eraser changes committed. Annotations changed: {annotations_changed}\")\n #print(f\"Current annotations: {self.annotations}\")\n \n def discard_eraser_changes(self):\n self.temp_eraser_mask = None\n self.update()\n \n \n def paintEvent(self, event):\n super().paintEvent(event)\n if self.scaled_pixmap:\n painter = QPainter(self)\n painter.setRenderHint(QPainter.Antialiasing)\n \n # Draw the image\n painter.drawPixmap(int(self.offset_x), int(self.offset_y), self.scaled_pixmap)\n \n # Draw annotations\n self.draw_annotations(painter)\n \n # Draw other elements\n if self.editing_polygon:\n self.draw_editing_polygon(painter)\n \n if self.drawing_rectangle and self.current_rectangle:\n self.draw_current_rectangle(painter)\n \n if self.sam_magic_wand_active and self.sam_bbox:\n self.draw_sam_bbox(painter)\n \n # Draw temporary paint mask\n if self.temp_paint_mask is not None:\n self.draw_temp_paint_mask(painter)\n \n # Draw temporary eraser mask\n if self.temp_eraser_mask is not None:\n self.draw_temp_eraser_mask(painter)\n \n # Draw brush/eraser size indicator\n self.draw_tool_size_indicator(painter)\n \n # Draw temporary YOLO predictions\n if self.temp_annotations:\n self.draw_temp_annotations(painter)\n \n painter.end()\n\n def draw_temp_annotations(self, painter):\n painter.save()\n painter.translate(self.offset_x, self.offset_y)\n painter.scale(self.zoom_factor, self.zoom_factor)\n\n for annotation in self.temp_annotations:\n color = QColor(255, 165, 0, 128) # Semi-transparent orange\n painter.setPen(QPen(color, 2 / self.zoom_factor, Qt.DashLine))\n painter.setBrush(QBrush(color))\n\n if \"bbox\" in annotation:\n x, y, w, h = annotation[\"bbox\"]\n painter.drawRect(QRectF(x, y, w, h))\n elif \"segmentation\" in annotation:\n points = [QPointF(float(x), float(y)) for x, y in zip(annotation[\"segmentation\"][0::2], annotation[\"segmentation\"][1::2])]\n painter.drawPolygon(QPolygonF(points))\n\n # Draw label and score\n painter.setFont(QFont(\"Arial\", int(12 / self.zoom_factor)))\n label = f\"{annotation['category_name']} {annotation['score']:.2f}\"\n if \"bbox\" in annotation:\n x, y, _, _ = annotation[\"bbox\"]\n painter.drawText(QPointF(x, y - 5), label)\n elif \"segmentation\" in annotation:\n centroid = self.calculate_centroid(points)\n if centroid:\n painter.drawText(centroid, label)\n\n painter.restore()\n \n def accept_temp_annotations(self):\n for annotation in self.temp_annotations:\n class_name = annotation['category_name']\n \n # Check if the class exists, if not, add it\n if class_name not in self.main_window.class_mapping:\n self.main_window.add_class(class_name)\n \n if class_name not in self.annotations:\n self.annotations[class_name] = []\n \n del annotation['temp']\n del annotation['score'] # Remove the score as it's not needed in the final annotation\n self.annotations[class_name].append(annotation)\n self.main_window.add_annotation_to_list(annotation)\n \n self.temp_annotations.clear()\n self.main_window.save_current_annotations()\n self.main_window.update_slice_list_colors()\n self.update()\n\n def discard_temp_annotations(self):\n self.temp_annotations.clear()\n self.update()\n \n def draw_temp_paint_mask(self, painter):\n if self.temp_paint_mask is not None:\n painter.save()\n painter.translate(self.offset_x, self.offset_y)\n painter.scale(self.zoom_factor, self.zoom_factor)\n \n mask_image = QImage(self.temp_paint_mask.data, self.temp_paint_mask.shape[1], self.temp_paint_mask.shape[0], self.temp_paint_mask.shape[1], QImage.Format_Grayscale8)\n mask_pixmap = QPixmap.fromImage(mask_image)\n painter.setOpacity(0.5)\n painter.drawPixmap(0, 0, mask_pixmap)\n painter.setOpacity(1.0)\n \n painter.restore()\n \n \n \n \n def draw_temp_eraser_mask(self, painter):\n if self.temp_eraser_mask is not None:\n painter.save()\n painter.translate(self.offset_x, self.offset_y)\n painter.scale(self.zoom_factor, self.zoom_factor)\n \n mask_image = QImage(self.temp_eraser_mask.data, self.temp_eraser_mask.shape[1], self.temp_eraser_mask.shape[0], self.temp_eraser_mask.shape[1], QImage.Format_Grayscale8)\n mask_pixmap = QPixmap.fromImage(mask_image)\n painter.setOpacity(0.5)\n painter.drawPixmap(0, 0, mask_pixmap)\n painter.setOpacity(1.0)\n \n painter.restore()\n \n \n\n \n\n def draw_tool_size_indicator(self, painter):\n if self.current_tool in [\"paint_brush\", \"eraser\"] and hasattr(self, 'cursor_pos'):\n painter.save()\n painter.translate(self.offset_x, self.offset_y)\n painter.scale(self.zoom_factor, self.zoom_factor)\n \n if self.current_tool == \"paint_brush\":\n size = self.main_window.paint_brush_size\n color = QColor(255, 0, 0, 128) # Semi-transparent red\n else: # eraser\n size = self.main_window.eraser_size\n color = QColor(0, 0, 255, 128) # Semi-transparent blue\n \n # Draw filled circle with lower opacity\n painter.setOpacity(0.3)\n painter.setPen(Qt.NoPen)\n painter.setBrush(color)\n painter.drawEllipse(QPointF(self.cursor_pos[0], self.cursor_pos[1]), size, size)\n \n # Draw circle outline with full opacity\n painter.setOpacity(1.0)\n painter.setPen(QPen(color.darker(150), 1 / self.zoom_factor, Qt.SolidLine))\n painter.setBrush(Qt.NoBrush)\n painter.drawEllipse(QPointF(self.cursor_pos[0], self.cursor_pos[1]), size, size)\n \n # Draw size text\n # Reset the transform to ensure text is drawn at screen coordinates\n painter.resetTransform()\n font = QFont()\n font.setPointSize(10)\n painter.setFont(font)\n painter.setPen(QPen(Qt.black)) # Use black color for better visibility\n \n # Convert cursor position back to screen coordinates\n screen_x = self.cursor_pos[0] * self.zoom_factor + self.offset_x\n screen_y = self.cursor_pos[1] * self.zoom_factor + self.offset_y\n \n # Position text above the circle\n text_rect = QRectF(screen_x + (size * self.zoom_factor), \n screen_y - (size * self.zoom_factor),\n 100, 20)\n \n text = f\"Size: {size}\"\n painter.drawText(text_rect, Qt.AlignLeft | Qt.AlignVCenter, text)\n \n painter.restore()\n\n\n def draw_paint_mask(self, painter):\n if self.paint_mask is not None:\n mask_image = QImage(self.paint_mask.data, self.paint_mask.shape[1], self.paint_mask.shape[0], self.paint_mask.shape[1], QImage.Format_Grayscale8)\n mask_pixmap = QPixmap.fromImage(mask_image)\n painter.setOpacity(0.5)\n painter.drawPixmap(self.offset_x, self.offset_y, mask_pixmap.scaled(self.scaled_pixmap.size()))\n painter.setOpacity(1.0)\n \n def draw_eraser_mask(self, painter):\n if self.eraser_mask is not None:\n mask_image = QImage(self.eraser_mask.data, self.eraser_mask.shape[1], self.eraser_mask.shape[0], self.eraser_mask.shape[1], QImage.Format_Grayscale8)\n mask_pixmap = QPixmap.fromImage(mask_image)\n painter.setOpacity(0.5)\n painter.drawPixmap(self.offset_x, self.offset_y, mask_pixmap.scaled(self.scaled_pixmap.size()))\n painter.setOpacity(1.0)\n \n\n def draw_sam_bbox(self, painter):\n painter.save()\n painter.translate(self.offset_x, self.offset_y)\n painter.scale(self.zoom_factor, self.zoom_factor)\n painter.setPen(QPen(Qt.red, 2 / self.zoom_factor, Qt.SolidLine))\n x1, y1, x2, y2 = self.sam_bbox\n painter.drawRect(QRectF(min(x1, x2), min(y1, y2), abs(x2 - x1), abs(y2 - y1)))\n painter.restore()\n \n def clear_temp_sam_prediction(self):\n self.temp_sam_prediction = None\n self.update()\n\n def check_unsaved_changes(self):\n if self.temp_paint_mask is not None or self.temp_eraser_mask is not None:\n reply = QMessageBox.question(\n self.main_window, 'Unsaved Changes',\n \"You have unsaved changes. Do you want to save them?\",\n QMessageBox.Yes | QMessageBox.No | QMessageBox.Cancel\n )\n if reply == QMessageBox.Yes:\n if self.temp_paint_mask is not None:\n self.commit_paint_annotation()\n if self.temp_eraser_mask is not None:\n self.commit_eraser_changes()\n return True\n elif reply == QMessageBox.No:\n self.discard_paint_annotation()\n self.discard_eraser_changes()\n return True\n else: # Cancel\n return False\n return True # No unsaved changes\n \n def clear(self):\n super().clear()\n self.annotations.clear()\n self.current_annotation.clear()\n self.temp_point = None\n self.current_tool = None\n self.start_point = None\n self.end_point = None\n self.highlighted_annotations.clear()\n self.original_pixmap = None\n self.scaled_pixmap = None\n self.editing_polygon = None\n self.editing_point_index = None\n self.hover_point_index = None\n self.current_rectangle = None\n self.sam_bbox = None\n self.temp_sam_prediction = None\n self.update()\n\n\n def set_class_visibility(self, class_name, is_visible):\n self.class_visibility[class_name] = is_visible\n\n def draw_annotations(self, painter):\n \"\"\"Draw all annotations on the image.\"\"\"\n if not self.original_pixmap:\n return\n \n painter.save()\n painter.translate(self.offset_x, self.offset_y)\n painter.scale(self.zoom_factor, self.zoom_factor)\n \n for class_name, class_annotations in self.annotations.items():\n if not self.main_window.is_class_visible(class_name):\n continue \n \n color = self.class_colors.get(class_name, QColor(Qt.white))\n for annotation in class_annotations:\n if annotation in self.highlighted_annotations:\n border_color = Qt.red\n fill_color = QColor(Qt.red)\n else:\n border_color = color\n fill_color = QColor(color)\n \n fill_color.setAlphaF(self.fill_opacity)\n \n text_color = Qt.white if self.dark_mode else Qt.black\n painter.setPen(QPen(border_color, 2 / self.zoom_factor, Qt.SolidLine))\n painter.setBrush(QBrush(fill_color))\n \n if \"segmentation\" in annotation and annotation[\"segmentation\"] is not None:\n segmentation = annotation[\"segmentation\"]\n if isinstance(segmentation, list) and len(segmentation) > 0:\n if isinstance(segmentation[0], list): # Multiple polygons\n for polygon in segmentation:\n points = [QPointF(float(x), float(y)) for x, y in zip(polygon[0::2], polygon[1::2])]\n if points:\n painter.drawPolygon(QPolygonF(points))\n else: # Single polygon\n points = [QPointF(float(x), float(y)) for x, y in zip(segmentation[0::2], segmentation[1::2])]\n if points:\n painter.drawPolygon(QPolygonF(points))\n \n # Draw centroid and label\n if points:\n centroid = self.calculate_centroid(points)\n if centroid:\n painter.setFont(QFont(\"Arial\", int(12 / self.zoom_factor)))\n painter.setPen(QPen(text_color, 2 / self.zoom_factor, Qt.SolidLine))\n painter.drawText(centroid, f\"{class_name} {annotation.get('number', '')}\")\n \n elif \"bbox\" in annotation:\n x, y, width, height = annotation[\"bbox\"]\n painter.drawRect(QRectF(x, y, width, height))\n painter.setPen(QPen(text_color, 2 / self.zoom_factor, Qt.SolidLine))\n painter.drawText(QPointF(x, y), f\"{class_name} {annotation.get('number', '')}\")\n \n if self.current_annotation:\n painter.setPen(QPen(Qt.red, 2 / self.zoom_factor, Qt.SolidLine))\n points = [QPointF(float(x), float(y)) for x, y in self.current_annotation]\n if len(points) > 1:\n painter.drawPolyline(QPolygonF(points))\n for point in points:\n painter.drawEllipse(point, 5 / self.zoom_factor, 5 / self.zoom_factor)\n if self.temp_point:\n painter.drawLine(points[-1], QPointF(float(self.temp_point[0]), float(self.temp_point[1])))\n \n # Draw temporary SAM prediction\n if self.temp_sam_prediction:\n temp_color = QColor(255, 165, 0, 128) # Semi-transparent orange\n painter.setPen(QPen(temp_color, 2 / self.zoom_factor, Qt.DashLine))\n painter.setBrush(QBrush(temp_color))\n \n segmentation = self.temp_sam_prediction[\"segmentation\"]\n points = [QPointF(float(x), float(y)) for x, y in zip(segmentation[0::2], segmentation[1::2])]\n if points:\n painter.drawPolygon(QPolygonF(points))\n centroid = self.calculate_centroid(points)\n if centroid:\n painter.setFont(QFont(\"Arial\", int(12 / self.zoom_factor)))\n painter.drawText(centroid, f\"SAM: {self.temp_sam_prediction['score']:.2f}\")\n \n painter.restore()\n\n def draw_current_rectangle(self, painter):\n \"\"\"Draw the current rectangle being created.\"\"\"\n if not self.current_rectangle:\n return\n \n painter.save()\n painter.translate(self.offset_x, self.offset_y)\n painter.scale(self.zoom_factor, self.zoom_factor)\n\n x1, y1, x2, y2 = self.current_rectangle\n color = self.class_colors.get(self.main_window.current_class, QColor(Qt.red))\n painter.setPen(QPen(color, 2 / self.zoom_factor, Qt.SolidLine))\n painter.drawRect(QRectF(float(x1), float(y1), float(x2 - x1), float(y2 - y1)))\n\n painter.restore()\n\n def get_rectangle_from_points(self):\n \"\"\"Get rectangle coordinates from start and end points.\"\"\"\n if not self.start_point or not self.end_point:\n return None\n x1, y1 = self.start_point\n x2, y2 = self.end_point\n return [min(x1, x2), min(y1, y2), max(x1, x2), max(y1, y2)]\n\n def draw_editing_polygon(self, painter):\n \"\"\"Draw the polygon being edited.\"\"\"\n painter.save()\n painter.translate(self.offset_x, self.offset_y)\n painter.scale(self.zoom_factor, self.zoom_factor)\n\n points = [QPointF(float(x), float(y)) for x, y in zip(self.editing_polygon[\"segmentation\"][0::2], self.editing_polygon[\"segmentation\"][1::2])]\n color = self.class_colors.get(self.editing_polygon[\"category_name\"], QColor(Qt.white))\n fill_color = QColor(color)\n fill_color.setAlphaF(self.fill_opacity)\n \n painter.setPen(QPen(color, 2 / self.zoom_factor, Qt.SolidLine))\n painter.setBrush(QBrush(fill_color))\n painter.drawPolygon(QPolygonF(points)) # Changed QPolygon to QPolygonF - Sreeni\n\n for i, point in enumerate(points):\n if i == self.hover_point_index:\n painter.setBrush(QColor(255, 0, 0))\n else:\n painter.setBrush(QColor(0, 255, 0))\n painter.drawEllipse(point, 5 / self.zoom_factor, 5 / self.zoom_factor)\n\n painter.restore()\n\n def calculate_centroid(self, points):\n \"\"\"Calculate the centroid of a polygon.\"\"\"\n if not points:\n return None\n x_coords = [point.x() for point in points]\n y_coords = [point.y() for point in points]\n centroid_x = sum(x_coords) / len(points)\n centroid_y = sum(y_coords) / len(points)\n return QPointF(centroid_x, centroid_y)\n\n def set_zoom(self, zoom_factor):\n \"\"\"Set the zoom factor and update the display.\"\"\"\n self.zoom_factor = zoom_factor\n self.update_scaled_pixmap()\n self.update()\n\n def wheelEvent(self, event: QWheelEvent):\n if event.modifiers() == Qt.ControlModifier:\n delta = event.angleDelta().y()\n if delta > 0:\n self.main_window.zoom_in()\n else:\n self.main_window.zoom_out()\n event.accept()\n else:\n super().wheelEvent(event)\n\n def mousePressEvent(self, event: QMouseEvent):\n if not self.original_pixmap:\n return\n if event.modifiers() == Qt.ControlModifier and event.button() == Qt.LeftButton:\n self.pan_start_pos = event.pos()\n self.setCursor(Qt.ClosedHandCursor)\n event.accept()\n else:\n pos = self.get_image_coordinates(event.pos())\n if event.button() == Qt.LeftButton:\n if self.sam_magic_wand_active:\n self.sam_bbox = [pos[0], pos[1], pos[0], pos[1]]\n self.drawing_sam_bbox = True\n elif self.editing_polygon:\n self.handle_editing_click(pos, event)\n elif self.current_tool == \"polygon\":\n if not self.drawing_polygon:\n self.drawing_polygon = True\n self.current_annotation = []\n self.current_annotation.append(pos)\n elif self.current_tool == \"rectangle\":\n self.start_point = pos\n self.end_point = pos\n self.drawing_rectangle = True\n self.current_rectangle = None\n elif self.current_tool == \"paint_brush\":\n self.start_painting(pos)\n elif self.current_tool == \"eraser\":\n self.start_erasing(pos)\n self.update()\n\n def mouseMoveEvent(self, event: QMouseEvent):\n if not self.original_pixmap:\n return\n self.cursor_pos = self.get_image_coordinates(event.pos())\n \n if event.modifiers() == Qt.ControlModifier and event.buttons() == Qt.LeftButton:\n if self.pan_start_pos:\n delta = event.pos() - self.pan_start_pos\n scrollbar_h = self.main_window.scroll_area.horizontalScrollBar()\n scrollbar_v = self.main_window.scroll_area.verticalScrollBar()\n scrollbar_h.setValue(scrollbar_h.value() - delta.x())\n scrollbar_v.setValue(scrollbar_v.value() - delta.y())\n self.pan_start_pos = event.pos()\n event.accept()\n else:\n pos = self.cursor_pos\n if self.sam_magic_wand_active and self.drawing_sam_bbox:\n if self.sam_bbox is not None:\n self.sam_bbox[2] = pos[0]\n self.sam_bbox[3] = pos[1]\n elif self.editing_polygon:\n self.handle_editing_move(pos)\n elif self.current_tool == \"polygon\" and self.current_annotation:\n self.temp_point = pos\n elif self.current_tool == \"rectangle\" and self.drawing_rectangle:\n self.end_point = pos\n self.current_rectangle = self.get_rectangle_from_points()\n elif self.current_tool == \"paint_brush\" and event.buttons() == Qt.LeftButton:\n self.continue_painting(pos)\n elif self.current_tool == \"eraser\" and event.buttons() == Qt.LeftButton:\n self.continue_erasing(pos)\n self.update()\n\n def mouseReleaseEvent(self, event: QMouseEvent):\n if not self.original_pixmap:\n return\n if event.modifiers() == Qt.ControlModifier and event.button() == Qt.LeftButton:\n self.pan_start_pos = None\n self.setCursor(Qt.ArrowCursor)\n event.accept()\n else:\n pos = self.get_image_coordinates(event.pos())\n if event.button() == Qt.LeftButton:\n if self.sam_magic_wand_active and self.drawing_sam_bbox:\n if self.sam_bbox is not None:\n self.sam_bbox[2] = pos[0]\n self.sam_bbox[3] = pos[1]\n self.drawing_sam_bbox = False\n self.main_window.apply_sam_prediction()\n elif self.editing_polygon:\n self.editing_point_index = None\n elif self.current_tool == \"rectangle\" and self.drawing_rectangle:\n self.drawing_rectangle = False\n if self.current_rectangle:\n self.main_window.finish_rectangle()\n elif self.current_tool == \"paint_brush\":\n self.finish_painting()\n elif self.current_tool == \"eraser\":\n self.finish_erasing()\n self.update()\n \n\n def mouseDoubleClickEvent(self, event):\n if not self.pixmap():\n return\n pos = self.get_image_coordinates(event.pos())\n if event.button() == Qt.LeftButton:\n if self.drawing_polygon and len(self.current_annotation) > 2:\n self.finish_polygon()\n else:\n self.clear_current_annotation()\n annotation = self.start_polygon_edit(pos)\n if annotation:\n self.main_window.select_annotation_in_list(annotation)\n self.update()\n\n def get_image_coordinates(self, pos):\n if not self.scaled_pixmap:\n return (0, 0)\n x = (pos.x() - self.offset_x) / self.zoom_factor\n y = (pos.y() - self.offset_y) / self.zoom_factor\n return (int(x), int(y))\n\n def keyPressEvent(self, event: QKeyEvent):\n if event.key() == Qt.Key_Return or event.key() == Qt.Key_Enter:\n if self.temp_annotations:\n self.accept_temp_annotations()\n elif self.temp_sam_prediction:\n self.main_window.accept_sam_prediction()\n elif self.editing_polygon:\n self.editing_polygon = None\n self.editing_point_index = None\n self.hover_point_index = None\n self.main_window.enable_tools()\n self.main_window.update_annotation_list()\n elif self.current_tool == \"polygon\" and self.drawing_polygon:\n self.finish_polygon()\n elif self.current_tool == \"paint_brush\":\n self.commit_paint_annotation()\n elif self.current_tool == \"eraser\":\n self.commit_eraser_changes()\n else:\n self.finish_current_annotation()\n elif event.key() == Qt.Key_Escape:\n if self.temp_annotations:\n self.discard_temp_annotations()\n elif self.sam_magic_wand_active:\n self.sam_bbox = None\n self.clear_temp_sam_prediction()\n elif self.editing_polygon:\n self.editing_polygon = None\n self.editing_point_index = None\n self.hover_point_index = None\n self.main_window.enable_tools()\n elif self.current_tool == \"paint_brush\":\n self.discard_paint_annotation()\n elif self.current_tool == \"eraser\":\n self.discard_eraser_changes()\n else:\n self.cancel_current_annotation()\n \n elif event.key() == Qt.Key_Delete:\n if self.editing_polygon:\n self.main_window.delete_selected_annotations()\n self.editing_polygon = None\n self.editing_point_index = None\n self.hover_point_index = None\n self.main_window.enable_tools()\n self.update()\n \n elif event.key() == Qt.Key_Minus:\n if self.current_tool == \"paint_brush\":\n self.main_window.paint_brush_size = max(1, self.main_window.paint_brush_size - 1)\n print(f\"Paint brush size: {self.main_window.paint_brush_size}\")\n elif self.current_tool == \"eraser\":\n self.main_window.eraser_size = max(1, self.main_window.eraser_size - 1)\n print(f\"Eraser size: {self.main_window.eraser_size}\")\n elif event.key() == Qt.Key_Equal:\n if self.current_tool == \"paint_brush\":\n self.main_window.paint_brush_size += 1\n print(f\"Paint brush size: {self.main_window.paint_brush_size}\")\n elif self.current_tool == \"eraser\":\n self.main_window.eraser_size += 1\n print(f\"Eraser size: {self.main_window.eraser_size}\")\n self.update()\n\n\n def cancel_current_annotation(self):\n \"\"\"Cancel the current annotation being created.\"\"\"\n if self.current_tool == \"polygon\" and self.current_annotation:\n self.current_annotation = []\n self.temp_point = None\n self.drawing_polygon = False\n self.update()\n \n\n def finish_current_annotation(self):\n \"\"\"Finish the current annotation being created.\"\"\"\n if self.current_tool == \"polygon\" and len(self.current_annotation) > 2:\n if self.main_window:\n self.main_window.finish_polygon()\n \n def finish_polygon(self):\n \"\"\"Finish the current polygon annotation.\"\"\"\n if self.drawing_polygon and len(self.current_annotation) > 2:\n self.drawing_polygon = False\n if self.main_window:\n self.main_window.finish_polygon()\n\n\n def start_polygon_edit(self, pos):\n for class_name, annotations in self.annotations.items():\n for annotation in annotations:\n if \"segmentation\" in annotation:\n points = [QPoint(int(x), int(y)) for x, y in zip(annotation[\"segmentation\"][0::2], annotation[\"segmentation\"][1::2])]\n if self.point_in_polygon(pos, points):\n self.editing_polygon = annotation\n self.current_tool = None\n self.main_window.disable_tools()\n self.main_window.reset_tool_buttons()\n return annotation\n return None\n\n def handle_editing_click(self, pos, event):\n \"\"\"Handle clicks during polygon editing.\"\"\"\n points = [QPoint(int(x), int(y)) for x, y in zip(self.editing_polygon[\"segmentation\"][0::2], self.editing_polygon[\"segmentation\"][1::2])]\n for i, point in enumerate(points):\n if self.distance(pos, point) < 10 / self.zoom_factor:\n if event.modifiers() & Qt.ShiftModifier:\n # Delete point\n del self.editing_polygon[\"segmentation\"][i*2:i*2+2]\n else:\n # Start moving point\n self.editing_point_index = i\n return\n # Add new point\n for i in range(len(points)):\n if self.point_on_line(pos, points[i], points[(i+1) % len(points)]):\n self.editing_polygon[\"segmentation\"][i*2+2:i*2+2] = [pos[0], pos[1]]\n self.editing_point_index = i + 1\n return\n\n def handle_editing_move(self, pos):\n \"\"\"Handle mouse movement during polygon editing.\"\"\"\n points = [QPoint(int(x), int(y)) for x, y in zip(self.editing_polygon[\"segmentation\"][0::2], self.editing_polygon[\"segmentation\"][1::2])]\n self.hover_point_index = None\n for i, point in enumerate(points):\n if self.distance(pos, point) < 10 / self.zoom_factor:\n self.hover_point_index = i\n break\n if self.editing_point_index is not None:\n self.editing_polygon[\"segmentation\"][self.editing_point_index*2] = pos[0]\n self.editing_polygon[\"segmentation\"][self.editing_point_index*2+1] = pos[1]\n \n \n def exit_editing_mode(self):\n self.editing_polygon = None\n self.editing_point_index = None\n self.hover_point_index = None\n self.update()\n\n @staticmethod\n def point_in_polygon(point, polygon):\n \"\"\"Check if a point is inside a polygon.\"\"\"\n n = len(polygon)\n inside = False\n p1x, p1y = polygon[0].x(), polygon[0].y()\n for i in range(n + 1):\n p2x, p2y = polygon[i % n].x(), polygon[i % n].y()\n if point[1] > min(p1y, p2y):\n if point[1] <= max(p1y, p2y):\n if point[0] <= max(p1x, p2x):\n if p1y != p2y:\n xinters = (point[1] - p1y) * (p2x - p1x) / (p2y - p1y) + p1x\n if p1x == p2x or point[0] <= xinters:\n inside = not inside\n p1x, p1y = p2x, p2y\n return inside\n\n @staticmethod\n def point_to_tuple(point):\n \"\"\"Convert QPoint to tuple.\"\"\"\n if isinstance(point, QPoint):\n return (point.x(), point.y())\n return point\n\n @staticmethod\n def distance(p1, p2):\n \"\"\"Calculate distance between two points.\"\"\"\n p1 = ImageLabel.point_to_tuple(p1)\n p2 = ImageLabel.point_to_tuple(p2)\n return ((p1[0] - p2[0])**2 + (p1[1] - p2[1])**2)**0.5\n\n @staticmethod\n def point_on_line(p, start, end):\n \"\"\"Check if a point is on a line segment.\"\"\"\n p = ImageLabel.point_to_tuple(p)\n start = ImageLabel.point_to_tuple(start)\n end = ImageLabel.point_to_tuple(end)\n d1 = ImageLabel.distance(p, start)\n d2 = ImageLabel.distance(p, end)\n line_length = ImageLabel.distance(start, end)\n buffer = 0.1 # Adjust this value for more or less strict \"on-line\" detection\n return abs(d1 + d2 - line_length) < buffer\n"
},
{
"path": "src/digitalsreeni_image_annotator/image_patcher.py",
"content": "import os\nimport numpy as np\nfrom PyQt5.QtWidgets import (QDialog, QVBoxLayout, QHBoxLayout, QPushButton, QLabel, QFileDialog, \n QSpinBox, QProgressBar, QMessageBox, QListWidget, QDialogButtonBox,\n QGridLayout, QComboBox, QApplication, QScrollArea, QWidget)\n\n\nfrom PyQt5.QtCore import Qt, QThread, pyqtSignal\nfrom PyQt5.QtCore import QTimer, QEventLoop\nfrom tifffile import TiffFile, imsave\nfrom PIL import Image\nimport traceback\n\nclass DimensionDialog(QDialog):\n def __init__(self, shape, file_name, parent=None):\n super().__init__(parent)\n self.shape = shape\n self.file_name = file_name\n self.initUI()\n\n def initUI(self):\n layout = QVBoxLayout()\n self.setLayout(layout)\n\n layout.addWidget(QLabel(f\"File: {self.file_name}\"))\n layout.addWidget(QLabel(f\"Image shape: {self.shape}\"))\n layout.addWidget(QLabel(\"Assign dimensions:\"))\n\n grid_layout = QGridLayout()\n self.combos = []\n dimensions = ['T', 'Z', 'C', 'H', 'W']\n for i, dim in enumerate(self.shape):\n grid_layout.addWidget(QLabel(f\"Dimension {i} (size {dim}):\"), i, 0)\n combo = QComboBox()\n combo.addItems(dimensions)\n grid_layout.addWidget(combo, i, 1)\n self.combos.append(combo)\n layout.addLayout(grid_layout)\n\n self.button_box = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel)\n self.button_box.accepted.connect(self.accept)\n self.button_box.rejected.connect(self.reject)\n layout.addWidget(self.button_box)\n\n def get_dimensions(self):\n return [combo.currentText() for combo in self.combos]\n\nclass PatchingThread(QThread):\n progress = pyqtSignal(int)\n error = pyqtSignal(str)\n finished = pyqtSignal()\n dimension_required = pyqtSignal(object, str)\n\n\n def __init__(self, input_files, output_dir, patch_size, overlap, dimensions):\n super().__init__()\n self.input_files = input_files\n self.output_dir = output_dir\n self.patch_size = patch_size\n self.overlap = overlap # Changed to tuple (to handle overlap_x, overlap_y independently) - Sreeni\n self.dimensions = dimensions\n\n def run(self):\n try:\n total_files = len(self.input_files)\n for i, file_path in enumerate(self.input_files):\n self.patch_image(file_path)\n self.progress.emit(int((i + 1) / total_files * 100))\n self.finished.emit()\n except Exception as e:\n self.error.emit(str(e))\n traceback.print_exc()\n\n def patch_image(self, file_path):\n file_name = os.path.basename(file_path)\n file_name_without_ext, file_extension = os.path.splitext(file_name)\n\n if file_extension.lower() in ['.tif', '.tiff']:\n with TiffFile(file_path) as tif:\n images = tif.asarray()\n if images.ndim > 2:\n if file_path not in self.dimensions:\n self.dimension_required.emit(images.shape, file_name)\n self.wait()\n dimensions = self.dimensions.get(file_path)\n if dimensions:\n if 'H' in dimensions and 'W' in dimensions:\n h_index = dimensions.index('H')\n w_index = dimensions.index('W')\n for idx in np.ndindex(images.shape[:h_index] + images.shape[h_index+2:]):\n slice_idx = idx[:h_index] + (slice(None), slice(None)) + idx[h_index:]\n image = images[slice_idx]\n slice_name = '_'.join([f'{dim}{i+1}' for dim, i in zip(dimensions, idx) if dim not in ['H', 'W']])\n self.save_patches(image, f\"{file_name_without_ext}_{slice_name}\", file_extension)\n else:\n raise ValueError(\"You must assign both H and W dimensions.\")\n else:\n raise ValueError(\"Dimensions were not properly assigned.\")\n else:\n self.save_patches(images, file_name_without_ext, file_extension)\n else:\n with Image.open(file_path) as img:\n image = np.array(img)\n self.save_patches(image, file_name_without_ext, file_extension)\n\n def save_patches(self, image, base_name, extension):\n h, w = image.shape[:2]\n patch_h, patch_w = self.patch_size\n overlap_x, overlap_y = self.overlap\n\n for i in range(0, h - overlap_y, patch_h - overlap_y):\n for j in range(0, w - overlap_x, patch_w - overlap_x):\n if i + patch_h <= h and j + patch_w <= w: # Only save full-sized patches\n patch = image[i:i+patch_h, j:j+patch_w]\n patch_name = f\"{base_name}_patch_{i}_{j}{extension}\"\n output_path = os.path.join(self.output_dir, patch_name)\n\n if extension.lower() in ['.tif', '.tiff']:\n imsave(output_path, patch)\n else:\n Image.fromarray(patch).save(output_path)\n\nclass ImagePatcherTool(QDialog):\n def __init__(self, parent=None):\n super().__init__(parent)\n self.setWindowModality(Qt.ApplicationModal)\n self.dimensions = {}\n self.input_files = []\n self.output_dir = \"\"\n self.initUI()\n\n def initUI(self):\n layout = QVBoxLayout()\n self.setLayout(layout)\n \n # Input files selection\n input_layout = QHBoxLayout()\n self.input_label = QLabel(\"Input Files:\")\n self.input_button = QPushButton(\"Select Files\")\n self.input_button.clicked.connect(self.select_input_files)\n input_layout.addWidget(self.input_label)\n input_layout.addWidget(self.input_button)\n layout.addLayout(input_layout)\n \n # Output directory selection\n output_layout = QHBoxLayout()\n self.output_label = QLabel(\"Output Directory:\")\n self.output_button = QPushButton(\"Select Directory\")\n self.output_button.clicked.connect(self.select_output_directory)\n output_layout.addWidget(self.output_label)\n output_layout.addWidget(self.output_button)\n layout.addLayout(output_layout)\n \n # Patch size inputs\n patch_layout = QHBoxLayout()\n patch_layout.addWidget(QLabel(\"Patch Size (W x H):\"))\n self.patch_w = QSpinBox()\n self.patch_w.setRange(1, 10000)\n self.patch_w.setValue(256)\n self.patch_h = QSpinBox()\n self.patch_h.setRange(1, 10000)\n self.patch_h.setValue(256)\n patch_layout.addWidget(self.patch_w)\n patch_layout.addWidget(self.patch_h)\n layout.addLayout(patch_layout)\n \n # Overlap inputs\n overlap_layout = QHBoxLayout()\n overlap_layout.addWidget(QLabel(\"Overlap (X, Y):\"))\n self.overlap_x = QSpinBox()\n self.overlap_x.setRange(0, 1000)\n self.overlap_x.setValue(0)\n self.overlap_y = QSpinBox()\n self.overlap_y.setRange(0, 1000)\n self.overlap_y.setValue(0)\n overlap_layout.addWidget(self.overlap_x)\n overlap_layout.addWidget(self.overlap_y)\n layout.addLayout(overlap_layout)\n \n # Create a scroll area for patch info\n scroll_area = QScrollArea()\n scroll_area.setWidgetResizable(True)\n scroll_area.setMinimumHeight(200) # Set a minimum height for the scroll area\n \n # Create a widget to hold the patch info label\n self.patch_info_container = QWidget()\n patch_info_layout = QVBoxLayout(self.patch_info_container)\n \n # Add the patch info label to the container\n self.patch_info_label = QLabel()\n self.patch_info_label.setAlignment(Qt.AlignLeft | Qt.AlignTop)\n patch_info_layout.addWidget(self.patch_info_label)\n \n # Set the container as the scroll area's widget\n scroll_area.setWidget(self.patch_info_container)\n \n # Add the scroll area to the main layout\n layout.addWidget(scroll_area)\n \n # Start button\n self.start_button = QPushButton(\"Start Patching\")\n self.start_button.clicked.connect(self.start_patching)\n layout.addWidget(self.start_button)\n \n # Progress bar\n self.progress_bar = QProgressBar()\n layout.addWidget(self.progress_bar)\n \n self.setWindowTitle('Image Patcher Tool')\n self.setMinimumWidth(500) # Set a minimum width for the dialog\n self.setMinimumHeight(600) # Set a minimum height for the dialog\n \n # Connect value changed signals\n self.patch_w.valueChanged.connect(self.update_patch_info)\n self.patch_h.valueChanged.connect(self.update_patch_info)\n self.overlap_x.valueChanged.connect(self.update_patch_info)\n self.overlap_y.valueChanged.connect(self.update_patch_info)\n\n def select_input_files(self):\n file_dialog = QFileDialog()\n self.input_files, _ = file_dialog.getOpenFileNames(self, \"Select Input Files\", \"\", \"Image Files (*.png *.jpg *.bmp *.tif *.tiff)\")\n self.input_label.setText(f\"Input Files: {len(self.input_files)} selected\")\n QApplication.processEvents()\n self.process_tiff_files()\n self.update_patch_info()\n \n def process_tiff_files(self):\n for file_path in self.input_files:\n if file_path.lower().endswith(('.tif', '.tiff')):\n self.check_tiff_dimensions(file_path)\n QApplication.processEvents()\n \n \n \n def check_tiff_dimensions(self, file_path):\n with TiffFile(file_path) as tif:\n images = tif.asarray()\n if images.ndim > 2:\n file_name = os.path.basename(file_path)\n dialog = DimensionDialog(images.shape, file_name, self)\n dialog.setWindowModality(Qt.ApplicationModal)\n result = dialog.exec_()\n if result == QDialog.Accepted:\n dimensions = dialog.get_dimensions()\n if 'H' in dimensions and 'W' in dimensions:\n self.dimensions[file_path] = dimensions\n else:\n QMessageBox.warning(self, \"Invalid Dimensions\", f\"You must assign both H and W dimensions for {file_name}.\")\n QApplication.processEvents()\n\n\n\n def select_output_directory(self):\n file_dialog = QFileDialog()\n self.output_dir = file_dialog.getExistingDirectory(self, \"Select Output Directory\")\n dir_name = os.path.basename(self.output_dir) if self.output_dir else \"\"\n self.output_label.setText(f\"Output Directory: {dir_name}\")\n QApplication.processEvents()\n self.update_patch_info()\n \n \n def start_patching(self):\n if not self.input_files:\n QMessageBox.warning(self, \"No Input Files\", \"Please select input files.\")\n return\n if not self.output_dir:\n QMessageBox.warning(self, \"No Output Directory\", \"Please select an output directory.\")\n return\n\n patch_size = (self.patch_h.value(), self.patch_w.value())\n overlap = (self.overlap_x.value(), self.overlap_y.value())\n\n self.patching_thread = PatchingThread(self.input_files, self.output_dir, patch_size, overlap, self.dimensions)\n self.patching_thread.progress.connect(self.update_progress)\n self.patching_thread.error.connect(self.show_error)\n self.patching_thread.finished.connect(self.patching_finished)\n self.patching_thread.dimension_required.connect(self.get_dimensions)\n self.patching_thread.start()\n \n self.start_button.setEnabled(False)\n\n\n def get_dimensions(self, shape, file_name):\n dialog = DimensionDialog(shape, file_name, self)\n dialog.setWindowModality(Qt.ApplicationModal)\n result = dialog.exec_()\n \n if result == QDialog.Accepted:\n dimensions = dialog.get_dimensions()\n if 'H' in dimensions and 'W' in dimensions:\n self.dimensions[file_name] = dimensions\n else:\n QMessageBox.warning(self, \"Invalid Dimensions\", f\"You must assign both H and W dimensions for {file_name}.\")\n QApplication.processEvents()\n self.patching_thread.wake()\n\n\n\n def get_patch_info(self):\n patch_info = {}\n patch_w = self.patch_w.value()\n patch_h = self.patch_h.value()\n overlap_x = self.overlap_x.value()\n overlap_y = self.overlap_y.value()\n\n for file_path in self.input_files:\n file_name = os.path.basename(file_path)\n if file_path.lower().endswith(('.tif', '.tiff')):\n with TiffFile(file_path) as tif:\n images = tif.asarray()\n if images.ndim > 2:\n dimensions = self.dimensions.get(file_path)\n if dimensions:\n h_index = dimensions.index('H')\n w_index = dimensions.index('W')\n h, w = images.shape[h_index], images.shape[w_index]\n else:\n h, w = images.shape[-2], images.shape[-1]\n else:\n h, w = images.shape\n else:\n with Image.open(file_path) as img:\n w, h = img.size\n\n patches_x = (w - overlap_x) // (patch_w - overlap_x)\n patches_y = (h - overlap_y) // (patch_h - overlap_y)\n leftover_x = w - (patches_x * (patch_w - overlap_x) + overlap_x)\n leftover_y = h - (patches_y * (patch_h - overlap_y) + overlap_y)\n\n patch_info[file_name] = {\n 'patches_x': patches_x,\n 'patches_y': patches_y,\n 'leftover_x': leftover_x,\n 'leftover_y': leftover_y\n }\n\n return patch_info\n\n def update_patch_info(self):\n if not self.input_files:\n self.patch_info_label.setText(\"No input files selected\")\n return\n \n patch_info = self.get_patch_info()\n if patch_info:\n info_text = \"Patch Information:
\"\n for file_name, info in patch_info.items():\n info_text += f\"File: {file_name}
\"\n info_text += f\"Patches: X: {info['patches_x']}, Y: {info['patches_y']}
\"\n info_text += f\"Leftover pixels: X: {info['leftover_x']}, Y: {info['leftover_y']}
\"\n info_text += \"\"\n self.patch_info_label.setText(info_text)\n else:\n self.patch_info_label.setText(\"Unable to calculate patch information\")\n\n\n\n\n\n def update_progress(self, value):\n self.progress_bar.setValue(value)\n\n def show_error(self, error_message):\n QMessageBox.critical(self, \"Error\", f\"An error occurred during patching:\\n{error_message}\")\n self.start_button.setEnabled(True)\n\n def patching_finished(self):\n self.start_button.setEnabled(True)\n QMessageBox.information(self, \"Patching Complete\", \"Image patching has been completed.\")\n\n def show_centered(self, parent):\n parent_geo = parent.geometry()\n self.move(parent_geo.center() - self.rect().center())\n self.show()\n\ndef show_image_patcher(parent=None):\n dialog = ImagePatcherTool(parent)\n dialog.show_centered(parent)\n return dialog"
},
{
"path": "src/digitalsreeni_image_annotator/import_formats.py",
"content": "\nimport json\nimport os\nimport yaml\nfrom PIL import Image\n\nfrom PyQt5.QtCore import QRectF\nfrom PyQt5.QtGui import QColor\nfrom PyQt5.QtWidgets import QMessageBox, QFileDialog\n\nimport os\nimport json\nfrom PyQt5.QtWidgets import QMessageBox\n\ndef import_coco_json(file_path, class_mapping):\n try:\n with open(file_path, 'r') as f:\n coco_data = json.load(f)\n\n # Validate required fields\n required_fields = ['images', 'annotations', 'categories']\n for field in required_fields:\n if field not in coco_data:\n raise ValueError(f\"Missing required field '{field}' in JSON file\")\n\n imported_annotations = {}\n image_info = {}\n\n # Create reverse mapping of category IDs to names\n category_id_to_name = {cat['id']: cat['name'] for cat in coco_data['categories']}\n\n # Determine the image directory\n json_dir = os.path.dirname(file_path)\n images_dir = os.path.join(json_dir, 'images')\n \n if not os.path.exists(images_dir):\n print(f\"Warning: 'images' subdirectory not found at {images_dir}\")\n\n # Process images\n for image in coco_data['images']:\n try:\n file_name = image['file_name']\n image_path = os.path.join(images_dir, file_name)\n \n image_info[image['id']] = {\n 'file_name': file_name,\n 'width': int(image['width']), # Ensure integers\n 'height': int(image['height']),\n 'path': image_path,\n 'id': int(image['id'])\n }\n except KeyError as e:\n print(f\"Warning: Missing required field in image data: {e}\")\n continue\n\n # Process annotations\n # Process annotations\n for ann in coco_data['annotations']:\n try:\n image_id = int(ann['image_id'])\n if image_id not in image_info:\n print(f\"Warning: Annotation refers to non-existent image ID: {image_id}\")\n continue\n\n if ann['category_id'] not in category_id_to_name:\n print(f\"Warning: Invalid category ID: {ann['category_id']}\")\n continue\n\n file_name = image_info[image_id]['file_name']\n category_name = category_id_to_name[ann['category_id']]\n\n if file_name not in imported_annotations:\n imported_annotations[file_name] = {}\n\n if category_name not in imported_annotations[file_name]:\n imported_annotations[file_name][category_name] = []\n\n annotation = {\n 'category_id': int(ann['category_id']),\n 'category_name': category_name\n }\n\n # Handle segmentation data\n has_valid_segmentation = False\n if 'segmentation' in ann and ann['segmentation']: # Check if segmentation exists and is not empty\n seg_data = ann['segmentation']\n if isinstance(seg_data, list):\n if seg_data and isinstance(seg_data[0], list):\n # Take the first polygon if multiple are present\n annotation['segmentation'] = [float(x) for x in seg_data[0]]\n has_valid_segmentation = True\n elif seg_data: # Single polygon\n annotation['segmentation'] = [float(x) for x in seg_data]\n has_valid_segmentation = True\n\n # If no valid segmentation but bbox exists, create segmentation from bbox\n if not has_valid_segmentation and 'bbox' in ann:\n x, y, w, h = [float(x) for x in ann['bbox']]\n # Create rectangle polygon from bbox [x,y, x+w,y, x+w,y+h, x,y+h]\n annotation['segmentation'] = [x, y, x + w, y, x + w, y + h, x, y + h]\n annotation['type'] = 'polygon'\n # Also store bbox data\n annotation['bbox'] = [x, y, w, h]\n elif has_valid_segmentation:\n annotation['type'] = 'polygon'\n elif 'bbox' in ann: # Fallback to pure bbox if no segmentation could be created\n annotation['bbox'] = [float(x) for x in ann['bbox']]\n annotation['type'] = 'rectangle'\n\n imported_annotations[file_name][category_name].append(annotation)\n \n except (KeyError, ValueError, TypeError) as e:\n print(f\"Warning: Error processing annotation: {e}\")\n continue\n\n return imported_annotations, image_info\n\n except json.JSONDecodeError as e:\n raise ValueError(f\"Invalid JSON file: {e}\")\n except Exception as e:\n raise ValueError(f\"Error importing COCO JSON: {e}\")\n\n\ndef import_yolo_v4(yaml_file_path, class_mapping):\n if not os.path.exists(yaml_file_path):\n raise ValueError(\"The selected YAML file does not exist.\")\n \n directory_path = os.path.dirname(yaml_file_path)\n \n with open(yaml_file_path, 'r') as f:\n yaml_data = yaml.safe_load(f)\n \n class_names = yaml_data.get('names', [])\n if not class_names:\n raise ValueError(\"No class names found in the YAML file.\")\n \n train_dir = os.path.join(directory_path, 'train')\n if not os.path.exists(train_dir):\n raise ValueError(\"No 'train' subdirectory found in the YAML file's directory.\")\n \n imported_annotations = {}\n image_info = {}\n \n images_dir = os.path.join(train_dir, 'images')\n labels_dir = os.path.join(train_dir, 'labels')\n \n if not os.path.exists(images_dir) or not os.path.exists(labels_dir):\n raise ValueError(\"The 'train' directory must contain both 'images' and 'labels' subdirectories.\")\n \n missing_images = []\n missing_labels = []\n \n for label_file in os.listdir(labels_dir):\n if label_file.lower().endswith('.txt'):\n base_name = os.path.splitext(label_file)[0]\n img_file = None\n img_path = None\n \n # Check for various image formats\n for ext in ['.jpg', '.jpeg', '.png', '.tiff', '.bmp', '.gif']:\n potential_img_file = base_name + ext\n potential_img_path = os.path.join(images_dir, potential_img_file)\n if os.path.exists(potential_img_path):\n img_file = potential_img_file\n img_path = potential_img_path\n break\n \n if img_path is None:\n missing_images.append(base_name)\n continue\n \n with Image.open(img_path) as img:\n img_width, img_height = img.size\n \n image_id = len(image_info) + 1\n image_info[image_id] = {\n 'file_name': img_file,\n 'width': img_width,\n 'height': img_height,\n 'id': image_id,\n 'path': img_path\n }\n \n imported_annotations[img_file] = {}\n \n label_path = os.path.join(labels_dir, label_file)\n with open(label_path, 'r') as f:\n lines = f.readlines()\n \n for line in lines:\n parts = line.strip().split()\n if len(parts) >= 5:\n class_id = int(parts[0])\n if class_id >= len(class_names):\n print(f\"Warning: Class ID {class_id} in {label_file} is out of range. Skipping this annotation.\")\n continue\n class_name = class_names[class_id]\n \n if class_name not in imported_annotations[img_file]:\n imported_annotations[img_file][class_name] = []\n \n if len(parts) == 5: # bounding box format\n x_center, y_center, width, height = map(float, parts[1:5])\n x1 = (x_center - width/2) * img_width\n y1 = (y_center - height/2) * img_height\n x2 = (x_center + width/2) * img_width\n y2 = (y_center + height/2) * img_height\n \n annotation = {\n 'category_id': class_id,\n 'category_name': class_name,\n 'type': 'rectangle',\n 'bbox': [x1, y1, x2-x1, y2-y1]\n }\n else: # polygon format\n polygon = [float(coord) * (img_width if i % 2 == 0 else img_height) for i, coord in enumerate(parts[1:])]\n \n annotation = {\n 'category_id': class_id,\n 'category_name': class_name,\n 'type': 'polygon',\n 'segmentation': polygon\n }\n \n imported_annotations[img_file][class_name].append(annotation)\n \n # Check for images without labels\n for img_file in os.listdir(images_dir):\n base_name, ext = os.path.splitext(img_file)\n if ext.lower() in ['.jpg', '.jpeg', '.png', '.tiff', '.bmp', '.gif']:\n label_file = base_name + '.txt'\n if not os.path.exists(os.path.join(labels_dir, label_file)):\n missing_labels.append(img_file)\n \n if missing_images or missing_labels:\n message = \"The following issues were found:\\n\\n\"\n if missing_images:\n message += f\"Labels without corresponding images: {', '.join(missing_images)}\\n\\n\"\n if missing_labels:\n message += f\"Images without corresponding labels: {', '.join(missing_labels)}\\n\\n\"\n message += \"Do you want to continue importing the remaining data?\"\n \n reply = QMessageBox.question(None, \"Import Issues\", message, \n QMessageBox.Yes | QMessageBox.No, QMessageBox.No)\n \n if reply == QMessageBox.No:\n raise ValueError(\"Import cancelled due to missing files.\")\n \n return imported_annotations, image_info\n\n\ndef import_yolo_v5plus(yaml_file_path, class_mapping):\n \"\"\"\n Import annotations from YOLO v5+ format.\n Expected directory structure:\n root_dir/\n ├── data.yaml\n ├── images/\n │ ├── train/\n │ └── val/\n └── labels/\n ├── train/\n └── val/\n \"\"\"\n if not os.path.exists(yaml_file_path):\n raise ValueError(\"The selected YAML file does not exist.\")\n \n root_dir = os.path.dirname(yaml_file_path)\n \n with open(yaml_file_path, 'r') as f:\n yaml_data = yaml.safe_load(f)\n \n class_names = yaml_data.get('names', [])\n if not class_names:\n raise ValueError(\"No class names found in the YAML file.\")\n \n imported_annotations = {}\n image_info = {}\n \n # Process both train and val directories\n for split in ['train', 'val']:\n images_dir = os.path.join(root_dir, 'images', split)\n labels_dir = os.path.join(root_dir, 'labels', split)\n \n if not os.path.exists(images_dir) or not os.path.exists(labels_dir):\n print(f\"Warning: {split} directory not found, skipping\")\n continue\n \n for label_file in os.listdir(labels_dir):\n if label_file.lower().endswith('.txt'):\n base_name = os.path.splitext(label_file)[0]\n img_file = None\n img_path = None\n \n # Check for various image formats\n for ext in ['.jpg', '.jpeg', '.png', '.tiff', '.bmp', '.gif']:\n potential_img_file = base_name + ext\n potential_img_path = os.path.join(images_dir, potential_img_file)\n if os.path.exists(potential_img_path):\n img_file = potential_img_file\n img_path = potential_img_path\n break\n \n if img_path is None:\n print(f\"Warning: No image found for label {label_file}\")\n continue\n \n with Image.open(img_path) as img:\n img_width, img_height = img.size\n \n image_id = len(image_info) + 1\n image_info[image_id] = {\n 'file_name': img_file,\n 'width': img_width,\n 'height': img_height,\n 'id': image_id,\n 'path': img_path\n }\n \n imported_annotations[img_file] = {}\n \n label_path = os.path.join(labels_dir, label_file)\n with open(label_path, 'r') as f:\n lines = f.readlines()\n \n for line in lines:\n parts = line.strip().split()\n if len(parts) >= 5:\n class_id = int(parts[0])\n if class_id >= len(class_names):\n print(f\"Warning: Class ID {class_id} in {label_file} is out of range\")\n continue\n class_name = class_names[class_id]\n \n if class_name not in imported_annotations[img_file]:\n imported_annotations[img_file][class_name] = []\n \n if len(parts) == 5: # bounding box format\n x_center, y_center, width, height = map(float, parts[1:5])\n x1 = (x_center - width/2) * img_width\n y1 = (y_center - height/2) * img_height\n w = width * img_width\n h = height * img_height\n \n annotation = {\n 'category_id': class_id,\n 'category_name': class_name,\n 'type': 'rectangle',\n 'bbox': [x1, y1, w, h]\n }\n else: # polygon format\n polygon = []\n for i in range(1, len(parts), 2):\n x = float(parts[i]) * img_width\n y = float(parts[i+1]) * img_height\n polygon.extend([x, y])\n \n annotation = {\n 'category_id': class_id,\n 'category_name': class_name,\n 'type': 'polygon',\n 'segmentation': polygon\n }\n \n imported_annotations[img_file][class_name].append(annotation)\n \n return imported_annotations, image_info\n\n\n\ndef process_import_format(import_format, file_path, class_mapping):\n if import_format == \"COCO JSON\":\n return import_coco_json(file_path, class_mapping)\n elif import_format == \"YOLO (v4 and earlier)\":\n return import_yolo_v4(file_path, class_mapping) # Still using same function, just updated format name\n elif import_format == \"YOLO (v5+)\":\n return import_yolo_v5plus(file_path, class_mapping) # New format handling\n else:\n raise ValueError(f\"Unsupported import format: {import_format}\")\n\n\n"
},
{
"path": "src/digitalsreeni_image_annotator/main.py",
"content": "\"\"\"\nMain entry point for the Image Annotator application.\n\nThis module creates and runs the main application window.\n\n@DigitalSreeni\nDr. Sreenivas Bhattiprolu\n\"\"\"\n\nimport sys\nimport os\nfrom PyQt5.QtWidgets import QApplication\nfrom .annotator_window import ImageAnnotator\n\n# To address Linux errors, by removing the QT_QPA_PLATFORM_PLUGIN_PATH \n# environment variable on Linux systems, which allows the application \n# to use the system's Qt platform plugins instead of potentially conflicting ones\nif sys.platform.startswith(\"linux\"):\n os.environ.pop(\"QT_QPA_PLATFORM_PLUGIN_PATH\", None)\n\ndef main():\n \"\"\"\n Main function to run the Image Annotator application.\n \"\"\"\n app = QApplication(sys.argv)\n window = ImageAnnotator()\n window.show()\n sys.exit(app.exec_())\n\nif __name__ == \"__main__\":\n main()"
},
{
"path": "src/digitalsreeni_image_annotator/project_details.py",
"content": "from PyQt5.QtWidgets import (QDialog, QVBoxLayout, QTextEdit, QPushButton, QLabel, \n QDialogButtonBox, QScrollArea, QWidget)\nfrom PyQt5.QtCore import Qt\nfrom PyQt5.QtGui import QFont\nimport os\nfrom datetime import datetime\n\nclass ProjectDetailsDialog(QDialog):\n def __init__(self, parent=None, stats_dialog=None):\n super().__init__(parent)\n self.parent = parent\n self.stats_dialog = stats_dialog\n self.setWindowTitle(\"Project Details\")\n self.setModal(True)\n self.setMinimumSize(600, 800) # Set initial size\n self.original_notes = parent.project_notes if parent else \"\"\n self.setup_ui()\n\n def setup_ui(self):\n layout = QVBoxLayout(self)\n\n scroll_area = QScrollArea()\n scroll_area.setWidgetResizable(True)\n scroll_content = QWidget()\n scroll_layout = QVBoxLayout(scroll_content)\n\n # Helper function to create bold labels\n def bold_label(text):\n label = QLabel(text)\n font = label.font()\n font.setBold(True)\n label.setFont(font)\n return label\n\n # Helper function to format datetime\n def format_datetime(date_string):\n try:\n dt = datetime.fromisoformat(date_string)\n return dt.strftime(\"%Y-%m-%d %H:%M:%S\")\n except ValueError:\n return date_string # Return original string if parsing fails\n\n # Project metadata\n scroll_layout.addWidget(bold_label(\"Project:\"))\n scroll_layout.addWidget(QLabel(os.path.basename(self.parent.current_project_file)))\n scroll_layout.addWidget(bold_label(\"Creation Date:\"))\n scroll_layout.addWidget(QLabel(format_datetime(getattr(self.parent, 'project_creation_date', 'N/A'))))\n scroll_layout.addWidget(bold_label(\"Last Modified:\"))\n scroll_layout.addWidget(QLabel(format_datetime(getattr(self.parent, 'last_modified', 'N/A'))))\n\n # Image information\n image_count = len(self.parent.all_images)\n scroll_layout.addWidget(bold_label(f\"Total Images: {image_count}\"))\n \n # List image file names\n scroll_layout.addWidget(bold_label(\"Image Files:\"))\n image_names = [f\"• {os.path.basename(path)}\" for path in self.parent.image_paths.values()]\n image_list = QLabel(\"\\n\".join(image_names))\n image_list.setWordWrap(True)\n scroll_layout.addWidget(image_list)\n\n # Multi-dimensional image information\n multi_slice_images = [img for img in self.parent.all_images if img.get('is_multi_slice', False)]\n if multi_slice_images:\n scroll_layout.addWidget(bold_label(f\"Multi-dimensional Images: {len(multi_slice_images)}\"))\n for img in multi_slice_images:\n slice_count = len(img.get('slices', []))\n scroll_layout.addWidget(QLabel(f\"• {os.path.basename(img['file_name'])}: {slice_count} slices\"))\n\n # Annotation information\n class_names = list(self.parent.class_mapping.keys())\n scroll_layout.addWidget(bold_label(\"Classes:\"))\n class_list = QLabel(\"\\n\".join([f\"• {name}\" for name in class_names]))\n class_list.setWordWrap(True)\n scroll_layout.addWidget(class_list)\n\n # Add annotation statistics\n if self.stats_dialog:\n scroll_layout.addWidget(bold_label(\"Annotation Statistics:\"))\n stats_text = self.stats_dialog.text_browser.toPlainText()\n stats_lines = stats_text.split('\\n')\n formatted_stats = []\n for line in stats_lines:\n if ':' in line:\n key, value = line.split(':', 1)\n formatted_stats.append(f\"{key}:{value}
\")\n else:\n formatted_stats.append(f\"{line}
\")\n stats_label = QLabel(\"\".join(formatted_stats))\n stats_label.setTextFormat(Qt.RichText)\n stats_label.setWordWrap(True)\n scroll_layout.addWidget(stats_label)\n\n scroll_area.setWidget(scroll_content)\n layout.addWidget(scroll_area)\n\n # Project notes\n layout.addWidget(bold_label(\"Project Notes:\"))\n self.notes_edit = QTextEdit()\n self.notes_edit.setPlainText(getattr(self.parent, 'project_notes', ''))\n layout.addWidget(self.notes_edit)\n\n # Buttons\n button_box = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel)\n button_box.accepted.connect(self.accept)\n button_box.rejected.connect(self.reject)\n layout.addWidget(button_box)\n\n def get_notes(self):\n return self.notes_edit.toPlainText()\n\n def were_changes_made(self):\n return self.get_notes() != self.original_notes"
},
{
"path": "src/digitalsreeni_image_annotator/project_search.py",
"content": "from PyQt5.QtWidgets import (QDialog, QVBoxLayout, QHBoxLayout, QLineEdit, QPushButton, \n QDateEdit, QLabel, QListWidget, QDialogButtonBox, QFormLayout,\n QFileDialog, QMessageBox)\nfrom PyQt5.QtCore import Qt, QDate\nimport os\nimport json\nfrom datetime import datetime\n\nclass ProjectSearchDialog(QDialog):\n def __init__(self, parent=None):\n super().__init__(parent)\n self.parent = parent\n self.setWindowTitle(\"Search Projects\")\n self.setModal(True)\n self.setMinimumSize(600, 400)\n self.search_directory = \"\"\n self.setup_ui()\n\n def setup_ui(self):\n layout = QVBoxLayout(self)\n\n # Search criteria\n form_layout = QFormLayout()\n self.keyword_edit = QLineEdit()\n self.keyword_edit.setPlaceholderText(\"Enter search query (e.g., monkey AND dog AND (project_animals OR project_zoo))\")\n form_layout.addRow(\"Search Query:\", self.keyword_edit)\n\n self.start_date = QDateEdit()\n self.start_date.setCalendarPopup(True)\n self.start_date.setDate(QDate.currentDate().addYears(-1))\n form_layout.addRow(\"Start Date:\", self.start_date)\n\n self.end_date = QDateEdit()\n self.end_date.setCalendarPopup(True)\n self.end_date.setDate(QDate.currentDate())\n form_layout.addRow(\"End Date:\", self.end_date)\n\n layout.addLayout(form_layout)\n\n # Directory selection\n dir_layout = QHBoxLayout()\n self.dir_edit = QLineEdit()\n dir_layout.addWidget(self.dir_edit)\n dir_button = QPushButton(\"Browse\")\n dir_button.clicked.connect(self.browse_directory)\n dir_layout.addWidget(dir_button)\n layout.addLayout(dir_layout)\n\n # Search button\n search_button = QPushButton(\"Search\")\n search_button.clicked.connect(self.perform_search)\n layout.addWidget(search_button)\n\n # Results list\n self.results_list = QListWidget()\n self.results_list.itemDoubleClicked.connect(self.open_selected_project)\n layout.addWidget(self.results_list)\n\n # Buttons\n button_box = QDialogButtonBox(QDialogButtonBox.Close)\n button_box.rejected.connect(self.reject)\n layout.addWidget(button_box)\n\n def browse_directory(self):\n directory = QFileDialog.getExistingDirectory(self, \"Select Directory to Search\")\n if directory:\n self.search_directory = directory\n self.dir_edit.setText(directory)\n\n def perform_search(self):\n if not self.search_directory:\n QMessageBox.warning(self, \"No Directory\", \"Please select a directory to search.\")\n return\n\n query = self.keyword_edit.text()\n start_date = self.start_date.date().toPyDate()\n end_date = self.end_date.date().toPyDate()\n\n self.results_list.clear()\n\n for root, dirs, files in os.walk(self.search_directory):\n for filename in files:\n if filename.endswith('.iap'):\n project_path = os.path.join(root, filename)\n try:\n with open(project_path, 'r') as f:\n project_data = json.load(f)\n \n if self.project_matches(project_data, query, start_date, end_date):\n self.results_list.addItem(project_path)\n except Exception as e:\n print(f\"Error reading project file {filename}: {str(e)}\")\n\n if self.results_list.count() == 0:\n QMessageBox.information(self, \"Search Results\", \"No matching projects found.\")\n else:\n QMessageBox.information(self, \"Search Results\", f\"{self.results_list.count()} matching projects found.\")\n\n def project_matches(self, project_data, query, start_date, end_date):\n # Check date range\n creation_date = project_data.get('creation_date', '')\n if creation_date:\n try:\n creation_date = datetime.fromisoformat(creation_date).date()\n if creation_date < start_date or creation_date > end_date:\n return False\n except ValueError:\n print(f\"Invalid date format in project: {creation_date}\")\n\n if not query:\n return True\n\n return self.evaluate_query(query.lower(), project_data)\n\n def term_matches(self, term, project_data):\n # Search in project name\n if term in os.path.basename(project_data.get('current_project_file', '')).lower():\n return True\n \n # Search in classes\n if any(term in class_info['name'].lower() for class_info in project_data.get('classes', [])):\n return True\n \n # Search in image names\n if any(term in img['file_name'].lower() for img in project_data.get('images', [])):\n return True\n \n # Search in project notes\n if term in project_data.get('notes', '').lower():\n return True\n \n return False\n\n\n def evaluate_query(self, query, project_data):\n tokens = self.tokenize_query(query)\n return self.evaluate_tokens(tokens, project_data)\n\n def tokenize_query(self, query):\n tokens = []\n current_token = \"\"\n for char in query:\n if char in '()':\n if current_token:\n tokens.append(current_token)\n current_token = \"\"\n tokens.append(char)\n elif char.isspace():\n if current_token:\n tokens.append(current_token)\n current_token = \"\"\n else:\n current_token += char\n if current_token:\n tokens.append(current_token)\n return tokens\n\n def evaluate_tokens(self, tokens, project_data):\n def evaluate_expression():\n nonlocal i\n result = True\n current_op = 'and'\n\n while i < len(tokens):\n if tokens[i] == '(':\n i += 1\n sub_result = evaluate_expression()\n if current_op == 'and':\n result = result and sub_result\n else:\n result = result or sub_result\n elif tokens[i] == ')':\n return result\n elif tokens[i].lower() in ['and', 'or']:\n current_op = tokens[i].lower()\n else:\n term_result = self.term_matches(tokens[i], project_data)\n if current_op == 'and':\n result = result and term_result\n else:\n result = result or term_result\n i += 1\n return result\n\n i = 0\n return evaluate_expression()\n\n\n def keyword_matches(self, keyword, project_data):\n # Search in project name\n if keyword in os.path.basename(project_data.get('current_project_file', '')).lower().split():\n return True\n \n # Search in classes\n if any(keyword in class_info['name'].lower().split() for class_info in project_data.get('classes', [])):\n return True\n \n # Search in image names\n if any(keyword in img['file_name'].lower().split() for img in project_data.get('images', [])):\n return True\n \n # Search in project notes\n if keyword in project_data.get('notes', '').lower().split():\n return True\n \n # Search in creation date and last modified date\n if keyword in project_data.get('creation_date', '').lower().split() or keyword in project_data.get('last_modified', '').lower().split():\n return True\n \n return False\n\n def open_selected_project(self, item):\n project_file = item.text()\n self.parent.open_specific_project(project_file)\n self.accept()\n\ndef show_project_search(parent):\n dialog = ProjectSearchDialog(parent)\n dialog.exec_()"
},
{
"path": "src/digitalsreeni_image_annotator/sam_utils.py",
"content": "import numpy as np\nfrom PyQt5.QtGui import QImage, QColor\nfrom ultralytics import SAM\n\nclass SAMUtils:\n def __init__(self):\n self.sam_models = {\n \"SAM 2 tiny\": \"sam2_t.pt\",\n \"SAM 2 small\": \"sam2_s.pt\",\n \"SAM 2 base\": \"sam2_b.pt\",\n \"SAM 2 large\": \"sam2_l.pt\",\n \"SAM 2.1 tiny\": \"sam2.1_t.pt\",\n \"SAM 2.1 small\": \"sam2.1_s.pt\",\n \"SAM 2.1 base\": \"sam2.1_b.pt\",\n \"SAM 2.1 large\": \"sam2.1_l.pt\",\n }\n self.current_sam_model = None\n self.sam_model = None\n\n def change_sam_model(self, model_name):\n if model_name != \"Pick a SAM Model\":\n self.current_sam_model = model_name\n self.sam_model = SAM(self.sam_models[self.current_sam_model])\n print(f\"Changed SAM model to: {model_name}\")\n else:\n self.current_sam_model = None\n self.sam_model = None\n print(\"SAM model unset\")\n\n def qimage_to_numpy(self, qimage):\n width = qimage.width()\n height = qimage.height()\n fmt = qimage.format()\n\n if fmt == QImage.Format_Grayscale16:\n buffer = qimage.constBits().asarray(height * width * 2)\n image = np.frombuffer(buffer, dtype=np.uint16).reshape((height, width))\n image_8bit = self.normalize_16bit_to_8bit(image)\n return np.stack((image_8bit,) * 3, axis=-1)\n \n elif fmt == QImage.Format_RGB16:\n buffer = qimage.constBits().asarray(height * width * 2)\n image = np.frombuffer(buffer, dtype=np.uint16).reshape((height, width))\n image_8bit = self.normalize_16bit_to_8bit(image)\n return np.stack((image_8bit,) * 3, axis=-1)\n\n elif fmt == QImage.Format_Grayscale8:\n buffer = qimage.constBits().asarray(height * width)\n image = np.frombuffer(buffer, dtype=np.uint8).reshape((height, width))\n return np.stack((image,) * 3, axis=-1)\n \n elif fmt in [QImage.Format_RGB32, QImage.Format_ARGB32, QImage.Format_ARGB32_Premultiplied]:\n buffer = qimage.constBits().asarray(height * width * 4)\n image = np.frombuffer(buffer, dtype=np.uint8).reshape((height, width, 4))\n return image[:, :, :3]\n \n elif fmt == QImage.Format_RGB888:\n buffer = qimage.constBits().asarray(height * width * 3)\n image = np.frombuffer(buffer, dtype=np.uint8).reshape((height, width, 3))\n return image\n \n elif fmt == QImage.Format_Indexed8:\n buffer = qimage.constBits().asarray(height * width)\n image = np.frombuffer(buffer, dtype=np.uint8).reshape((height, width))\n color_table = qimage.colorTable()\n rgb_image = np.zeros((height, width, 3), dtype=np.uint8)\n for y in range(height):\n for x in range(width):\n rgb_image[y, x] = QColor(color_table[image[y, x]]).getRgb()[:3]\n return rgb_image\n \n else:\n converted_image = qimage.convertToFormat(QImage.Format_RGB32)\n buffer = converted_image.constBits().asarray(height * width * 4)\n image = np.frombuffer(buffer, dtype=np.uint8).reshape((height, width, 4))\n return image[:, :, :3]\n\n def normalize_16bit_to_8bit(self, array):\n return ((array - array.min()) / (array.max() - array.min()) * 255).astype(np.uint8)\n\n def apply_sam_prediction(self, image, bbox):\n try:\n image_np = self.qimage_to_numpy(image)\n results = self.sam_model(image_np, bboxes=[bbox])\n mask = results[0].masks.data[0].cpu().numpy()\n\n if mask is not None:\n print(f\"Mask shape: {mask.shape}, Mask sum: {mask.sum()}\")\n contours = self.mask_to_polygon(mask)\n print(f\"Contours generated: {len(contours)} contour(s)\")\n\n if not contours:\n print(\"No valid contours found\")\n return None\n\n prediction = {\n \"segmentation\": contours[0],\n \"score\": float(results[0].boxes.conf[0])\n }\n return prediction\n else:\n print(\"Failed to generate mask\")\n return None\n except Exception as e:\n print(f\"Error in applying SAM prediction: {str(e)}\")\n import traceback\n traceback.print_exc()\n return None\n\n def mask_to_polygon(self, mask):\n import cv2\n contours, _ = cv2.findContours((mask > 0).astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)\n polygons = []\n for contour in contours:\n if cv2.contourArea(contour) > 10:\n polygon = contour.flatten().tolist()\n if len(polygon) >= 6:\n polygons.append(polygon)\n print(f\"Generated {len(polygons)} valid polygons\")\n return polygons\n"
},
{
"path": "src/digitalsreeni_image_annotator/slice_registration.py",
"content": "from PyQt5.QtWidgets import (QDialog, QVBoxLayout, QHBoxLayout, QPushButton, QFileDialog, \n QLabel, QComboBox, QMessageBox, QProgressDialog, QRadioButton,\n QButtonGroup, QSpinBox, QApplication, QGroupBox, QDoubleSpinBox)\nfrom PyQt5.QtCore import Qt\nfrom pystackreg import StackReg\nfrom skimage import io\nimport tifffile\nfrom PIL import Image\nimport numpy as np\nimport os\n\nclass SliceRegistrationTool(QDialog):\n def __init__(self, parent=None):\n super().__init__(parent)\n self.setWindowTitle(\"Slice Registration\")\n self.setGeometry(100, 100, 600, 400)\n self.setWindowFlags(self.windowFlags() | Qt.Window)\n self.setWindowModality(Qt.ApplicationModal) # Add modal behavior\n \n # Initialize variables first\n self.input_path = \"\"\n self.output_directory = \"\"\n \n self.initUI()\n\n def initUI(self):\n layout = QVBoxLayout()\n layout.setSpacing(10) # Add consistent spacing\n\n # Input selection\n input_group = QGroupBox(\"Input Selection\")\n input_layout = QVBoxLayout()\n \n self.dir_radio = QRadioButton(\"Directory of Image Files\")\n self.stack_radio = QRadioButton(\"TIFF Stack\")\n \n input_group = QButtonGroup(self)\n input_group.addButton(self.dir_radio)\n input_group.addButton(self.stack_radio)\n \n input_layout.addWidget(self.dir_radio)\n input_layout.addWidget(self.stack_radio)\n self.dir_radio.setChecked(True)\n \n # Input/Output file selection with labels\n self.input_label = QLabel(\"No input selected\")\n self.output_label = QLabel(\"No output directory selected\")\n \n file_select_layout = QVBoxLayout()\n \n input_file_layout = QHBoxLayout()\n self.select_input_btn = QPushButton(\"Select Input\")\n self.select_input_btn.clicked.connect(self.select_input)\n input_file_layout.addWidget(self.select_input_btn)\n input_file_layout.addWidget(self.input_label)\n \n output_file_layout = QHBoxLayout()\n self.select_output_btn = QPushButton(\"Select Output Directory\")\n self.select_output_btn.clicked.connect(self.select_output)\n output_file_layout.addWidget(self.select_output_btn)\n output_file_layout.addWidget(self.output_label)\n \n file_select_layout.addLayout(input_file_layout)\n file_select_layout.addLayout(output_file_layout)\n input_layout.addLayout(file_select_layout)\n \n layout.addLayout(input_layout)\n\n # Transform type\n transform_group = QGroupBox(\"Transformation Settings\")\n transform_layout = QVBoxLayout()\n \n transform_combo_layout = QHBoxLayout()\n transform_combo_layout.addWidget(QLabel(\"Type:\"))\n self.transform_combo = QComboBox()\n self.transform_combo.addItems([\n \"Translation (X-Y Translation Only)\",\n \"Rigid Body (Translation + Rotation)\",\n \"Scaled Rotation (Translation + Rotation + Scaling)\",\n \"Affine (Translation + Rotation + Scaling + Shearing)\",\n \"Bilinear (Non-linear; Does not preserve straight lines)\"\n ])\n transform_combo_layout.addWidget(self.transform_combo)\n transform_layout.addLayout(transform_combo_layout)\n transform_group.setLayout(transform_layout)\n layout.addWidget(transform_group)\n\n # Reference type\n ref_group = QGroupBox(\"Reference Settings\")\n ref_layout = QVBoxLayout()\n \n ref_combo_layout = QHBoxLayout()\n ref_combo_layout.addWidget(QLabel(\"Reference:\"))\n self.ref_combo = QComboBox()\n self.ref_combo.addItems([\n \"Previous Frame\",\n \"First Frame\",\n \"Mean of All Frames\",\n \"Mean of First N Frames\",\n \"Mean of First N Frames + Moving Average\"\n ])\n ref_combo_layout.addWidget(self.ref_combo)\n ref_layout.addLayout(ref_combo_layout)\n \n # N frames settings\n n_frames_layout = QHBoxLayout()\n n_frames_layout.addWidget(QLabel(\"N Frames:\"))\n self.n_frames_spin = QSpinBox()\n self.n_frames_spin.setRange(1, 100)\n self.n_frames_spin.setValue(10)\n self.n_frames_spin.setEnabled(False)\n n_frames_layout.addWidget(self.n_frames_spin)\n ref_layout.addLayout(n_frames_layout)\n \n # Moving average settings\n moving_avg_layout = QHBoxLayout()\n moving_avg_layout.addWidget(QLabel(\"Moving Average Window:\"))\n self.moving_avg_spin = QSpinBox()\n self.moving_avg_spin.setRange(1, 100)\n self.moving_avg_spin.setValue(10)\n self.moving_avg_spin.setEnabled(False)\n moving_avg_layout.addWidget(self.moving_avg_spin)\n ref_layout.addLayout(moving_avg_layout)\n \n ref_group.setLayout(ref_layout)\n layout.addWidget(ref_group)\n\n # Connect reference combo box\n self.ref_combo.currentTextChanged.connect(self.on_ref_changed)\n\n # Add spacing group\n spacing_group = QGroupBox(\"Pixel/Voxel Size\")\n spacing_layout = QVBoxLayout()\n \n # XY pixel size\n xy_size_layout = QHBoxLayout()\n xy_size_layout.addWidget(QLabel(\"XY Pixel Size:\"))\n self.xy_size_value = QDoubleSpinBox()\n self.xy_size_value.setRange(0.001, 1000.0)\n self.xy_size_value.setValue(1.0)\n self.xy_size_value.setDecimals(3)\n xy_size_layout.addWidget(self.xy_size_value)\n \n # Z spacing\n z_size_layout = QHBoxLayout()\n z_size_layout.addWidget(QLabel(\"Z Spacing:\"))\n self.z_size_value = QDoubleSpinBox()\n self.z_size_value.setRange(0.001, 1000.0)\n self.z_size_value.setValue(1.0)\n self.z_size_value.setDecimals(3)\n z_size_layout.addWidget(self.z_size_value)\n \n # Unit selector\n unit_layout = QHBoxLayout()\n unit_layout.addWidget(QLabel(\"Unit:\"))\n self.size_unit = QComboBox()\n self.size_unit.addItems([\"nm\", \"µm\", \"mm\"])\n self.size_unit.setCurrentText(\"µm\")\n unit_layout.addWidget(self.size_unit)\n \n spacing_layout.addLayout(xy_size_layout)\n spacing_layout.addLayout(z_size_layout)\n spacing_layout.addLayout(unit_layout)\n spacing_group.setLayout(spacing_layout)\n layout.addWidget(spacing_group)\n\n # Register button\n self.register_btn = QPushButton(\"Register\")\n self.register_btn.clicked.connect(self.register_slices)\n layout.addWidget(self.register_btn)\n\n self.setLayout(layout)\n\n def on_ref_changed(self, text):\n uses_n_frames = text in [\"Mean of First N Frames\", \"Mean of First N Frames + Moving Average\"]\n self.n_frames_spin.setEnabled(uses_n_frames)\n self.moving_avg_spin.setEnabled(text == \"Mean of First N Frames + Moving Average\")\n QApplication.processEvents() # Ensure UI updates\n \n\n def on_transform_changed(self, text):\n if text == \"Bilinear\" and self.ref_combo.currentText() == \"Previous\":\n QMessageBox.warning(self, \"Warning\", \n \"Bilinear transformation cannot be used with 'Previous' reference. \"\n \"Please select a different reference type.\")\n self.transform_combo.setCurrentText(\"Rigid Body\")\n\n\n def select_input(self):\n try:\n if self.dir_radio.isChecked():\n path = QFileDialog.getExistingDirectory(\n self,\n \"Select Directory with Images\",\n \"\",\n QFileDialog.ShowDirsOnly | QFileDialog.DontResolveSymlinks\n )\n else:\n path, _ = QFileDialog.getOpenFileName(\n self,\n \"Select TIFF Stack\",\n \"\",\n \"TIFF Files (*.tif *.tiff)\",\n options=QFileDialog.Options()\n )\n \n if path:\n self.input_path = path\n self.input_label.setText(f\"Selected: {os.path.basename(path)}\")\n self.input_label.setToolTip(path)\n QApplication.processEvents()\n \n except Exception as e:\n QMessageBox.critical(self, \"Error\", f\"Error selecting input: {str(e)}\")\n\n def select_output(self):\n try:\n directory = QFileDialog.getExistingDirectory(\n self,\n \"Select Output Directory\",\n \"\",\n QFileDialog.ShowDirsOnly | QFileDialog.DontResolveSymlinks\n )\n \n if directory:\n self.output_directory = directory\n self.output_label.setText(f\"Selected: {os.path.basename(directory)}\")\n self.output_label.setToolTip(directory)\n QApplication.processEvents()\n \n except Exception as e:\n QMessageBox.critical(self, \"Error\", f\"Error selecting output directory: {str(e)}\")\n\n def register_slices(self):\n if not self.input_path or not self.output_directory:\n QMessageBox.warning(self, \"Error\", \"Please select both input and output paths\")\n return\n \n try:\n progress = QProgressDialog(self)\n progress.setWindowTitle(\"Registration Progress\")\n progress.setLabelText(\"Loading images...\")\n progress.setMinimum(0)\n progress.setMaximum(100)\n progress.setWindowModality(Qt.WindowModal)\n progress.setMinimumWidth(400)\n progress.show()\n QApplication.processEvents()\n \n # Load images using scikit-image's imread\n if self.stack_radio.isChecked():\n progress.setLabelText(\"Loading TIFF stack...\")\n img0 = io.imread(self.input_path)\n else:\n progress.setLabelText(\"Loading images from directory...\")\n image_files = sorted([f for f in os.listdir(self.input_path) \n if f.lower().endswith(('.png', '.jpg', '.jpeg', '.tif', '.tiff'))])\n first_img = io.imread(os.path.join(self.input_path, image_files[0]))\n img0 = np.zeros((len(image_files), *first_img.shape), dtype=first_img.dtype)\n img0[0] = first_img\n for i, fname in enumerate(image_files[1:], 1):\n img0[i] = io.imread(os.path.join(self.input_path, fname))\n \n # Store original properties\n original_dtype = img0.dtype\n print(f\"Original image properties:\")\n print(f\"Dtype: {original_dtype}\")\n print(f\"Range: {img0.min()} - {img0.max()}\")\n print(f\"Shape: {img0.shape}\")\n \n progress.setValue(30)\n progress.setLabelText(\"Performing registration...\")\n QApplication.processEvents()\n \n # Set up StackReg with selected transformation\n transform_types = {\n \"Translation (X-Y Translation Only)\": StackReg.TRANSLATION,\n \"Rigid Body (Translation + Rotation)\": StackReg.RIGID_BODY,\n \"Scaled Rotation (Translation + Rotation + Scaling)\": StackReg.SCALED_ROTATION,\n \"Affine (Translation + Rotation + Scaling + Shearing)\": StackReg.AFFINE,\n \"Bilinear (Non-linear; Does not preserve straight lines)\": StackReg.BILINEAR\n }\n \n transform_type = transform_types[self.transform_combo.currentText()]\n sr = StackReg(transform_type)\n \n # Register images\n selected_ref = self.ref_combo.currentText()\n progress.setLabelText(f\"Registering images using {selected_ref}...\")\n progress.setValue(40)\n QApplication.processEvents()\n \n # Register and transform\n if selected_ref == \"Previous Frame\":\n out_registered = sr.register_transform_stack(img0, reference='previous')\n elif selected_ref == \"First Frame\":\n out_registered = sr.register_transform_stack(img0, reference='first')\n elif selected_ref == \"Mean of All Frames\":\n out_registered = sr.register_transform_stack(img0, reference='mean')\n elif selected_ref == \"Mean of First N Frames\":\n n_frames = self.n_frames_spin.value()\n out_registered = sr.register_transform_stack(img0, reference='first', n_frames=n_frames)\n elif selected_ref == \"Mean of First N Frames + Moving Average\":\n n_frames = self.n_frames_spin.value()\n moving_avg = self.moving_avg_spin.value()\n out_registered = sr.register_transform_stack(img0, reference='first', \n n_frames=n_frames, \n moving_average=moving_avg)\n \n progress.setValue(80)\n progress.setLabelText(\"Saving registered images...\")\n QApplication.processEvents()\n \n # Convert back to original dtype without changing values\n out_registered = out_registered.astype(original_dtype)\n \n print(f\"Output image properties:\")\n print(f\"Dtype: {out_registered.dtype}\")\n print(f\"Range: {out_registered.min()} - {out_registered.max()}\")\n print(f\"Shape: {out_registered.shape}\")\n \n # Save output\n if self.stack_radio.isChecked():\n output_name = os.path.splitext(os.path.basename(self.input_path))[0]\n else:\n output_name = \"registered_stack\"\n \n output_path = os.path.join(self.output_directory, f\"{output_name}_registered.tif\")\n \n # Get pixel sizes in micrometers (convert if necessary)\n xy_size = self.xy_size_value.value()\n z_size = self.z_size_value.value()\n unit = self.size_unit.currentText()\n \n # Convert to micrometers based on selected unit\n if unit == \"nm\":\n xy_size = xy_size / 1000\n z_size = z_size / 1000\n elif unit == \"mm\":\n xy_size = xy_size * 1000\n z_size = z_size * 1000\n \n # Save the stack\n tifffile.imwrite(\n output_path, \n out_registered,\n imagej=True,\n metadata={\n 'axes': 'ZYX',\n 'spacing': z_size, # Z spacing in micrometers\n 'unit': 'um',\n 'finterval': xy_size # XY pixel size in micrometers\n },\n resolution=(1.0/xy_size, 1.0/xy_size) # XY Resolution in pixels per unit\n )\n \n progress.setValue(100)\n QApplication.processEvents()\n \n # Include both XY and Z size info in success message\n QMessageBox.information(self, \"Success\", \n f\"Registration completed successfully!\\n\"\n f\"Output saved to:\\n{output_path}\\n\"\n f\"XY Pixel size: {self.xy_size_value.value()} {unit}\\n\"\n f\"Z Spacing: {self.z_size_value.value()} {unit}\")\n \n except Exception as e:\n print(f\"Error occurred: {str(e)}\")\n import traceback\n traceback.print_exc()\n QMessageBox.critical(self, \"Error\", str(e))\n \n \n def update_progress(self, progress_dialog, current_iteration, end_iteration):\n \"\"\"Helper function to update progress during registration\"\"\"\n if end_iteration > 0:\n percent = int(40 + (current_iteration / end_iteration) * 40) # Scale to 40-80% range\n progress_dialog.setValue(percent)\n progress_dialog.setLabelText(f\"Processing image {current_iteration}/{end_iteration}...\")\n QApplication.processEvents()\n \n\n def load_images(self):\n print(\"Starting image loading...\")\n try:\n if self.stack_radio.isChecked():\n print(f\"Loading TIFF stack from: {self.input_path}\")\n # Explicitly use scikit-image's imread for TIFF stacks\n stack = io.imread(self.input_path)\n if stack.dtype != np.float32:\n stack = stack.astype(np.float32)\n print(f\"Loaded TIFF stack shape: {stack.shape}\")\n return stack\n else:\n # Load individual images\n valid_extensions = ('.png', '.jpg', '.jpeg', '.bmp', '.tif', '.tiff')\n images = []\n files = sorted([f for f in os.listdir(self.input_path) \n if f.lower().endswith(valid_extensions)])\n \n print(f\"Found {len(files)} image files\")\n \n if not files:\n raise ValueError(\"No valid image files found in directory\")\n \n # Check first image size\n first_path = os.path.join(self.input_path, files[0])\n print(f\"Loading first image: {first_path}\")\n first_img = np.array(Image.open(first_path))\n ref_shape = first_img.shape\n images.append(first_img)\n print(f\"First image shape: {ref_shape}\")\n \n # Load remaining images and check sizes\n for f in files[1:]:\n img_path = os.path.join(self.input_path, f)\n print(f\"Loading: {f}\")\n img = np.array(Image.open(img_path))\n if img.shape != ref_shape:\n raise ValueError(f\"Image {f} has different dimensions from the first image\")\n images.append(img)\n \n stack = np.stack(images)\n print(f\"Final stack shape: {stack.shape}\")\n return stack\n \n except Exception as e:\n print(f\"Error in load_images: {str(e)}\")\n raise\n\n\n def show_centered(self, parent):\n parent_geo = parent.geometry()\n self.move(parent_geo.center() - self.rect().center())\n self.show()\n QApplication.processEvents() # Ensure window displays properly\n\n"
},
{
"path": "src/digitalsreeni_image_annotator/snake_game.py",
"content": "import sys\nimport random\nfrom PyQt5.QtWidgets import QApplication, QWidget, QDesktopWidget, QMessageBox\nfrom PyQt5.QtGui import QPainter, QColor\nfrom PyQt5.QtCore import Qt, QTimer\n\nclass SnakeGame(QWidget):\n def __init__(self):\n super().__init__()\n self.initUI()\n \n def initUI(self):\n self.setWindowTitle('Secret Snake Game')\n self.setFixedSize(600, 400) # Increased size\n self.center()\n \n self.snake = [(300, 200), (290, 200), (280, 200)]\n self.direction = 'RIGHT'\n self.food = self.place_food()\n self.score = 0\n \n self.timer = QTimer(self)\n self.timer.timeout.connect(self.update_game)\n self.timer.start(100)\n \n self.setFocusPolicy(Qt.StrongFocus)\n self.show()\n \n def center(self):\n qr = self.frameGeometry()\n cp = QDesktopWidget().availableGeometry().center()\n qr.moveCenter(cp)\n self.move(qr.topLeft())\n \n def paintEvent(self, event):\n painter = QPainter(self)\n painter.setRenderHint(QPainter.Antialiasing)\n \n # Draw snake\n painter.setBrush(QColor(0, 255, 0))\n for segment in self.snake:\n painter.drawRect(segment[0], segment[1], 10, 10)\n \n # Draw food\n painter.setBrush(QColor(255, 0, 0))\n painter.drawRect(self.food[0], self.food[1], 10, 10)\n \n # Draw score\n painter.setPen(QColor(0, 0, 0))\n painter.drawText(10, 20, f\"Score: {self.score}\")\n \n def keyPressEvent(self, event):\n key = event.key()\n \n if key == Qt.Key_Left and self.direction != 'RIGHT':\n self.direction = 'LEFT'\n elif key == Qt.Key_Right and self.direction != 'LEFT':\n self.direction = 'RIGHT'\n elif key == Qt.Key_Up and self.direction != 'DOWN':\n self.direction = 'UP'\n elif key == Qt.Key_Down and self.direction != 'UP':\n self.direction = 'DOWN'\n elif key == Qt.Key_Escape:\n self.close()\n \n def update_game(self):\n head = self.snake[0]\n \n if self.direction == 'LEFT':\n new_head = (head[0] - 10, head[1])\n elif self.direction == 'RIGHT':\n new_head = (head[0] + 10, head[1])\n elif self.direction == 'UP':\n new_head = (head[0], head[1] - 10)\n else: # DOWN\n new_head = (head[0], head[1] + 10)\n \n # Check if snake hit the edge\n if (new_head[0] < 0 or new_head[0] >= 600 or\n new_head[1] < 0 or new_head[1] >= 400):\n self.game_over()\n return\n \n self.snake.insert(0, new_head)\n \n if new_head == self.food:\n self.score += 1\n self.food = self.place_food()\n else:\n self.snake.pop()\n \n if new_head in self.snake[1:]:\n self.game_over()\n return\n \n self.update()\n \n def place_food(self):\n while True:\n x = random.randint(0, 59) * 10\n y = random.randint(0, 39) * 10\n if (x, y) not in self.snake:\n return (x, y)\n \n def game_over(self):\n self.timer.stop()\n QMessageBox.information(self, \"Game Over\", f\"Your score: {self.score}\")\n self.close()\n\nif __name__ == '__main__':\n app = QApplication(sys.argv)\n ex = SnakeGame()\n sys.exit(app.exec_())"
},
{
"path": "src/digitalsreeni_image_annotator/soft_dark_stylesheet.py",
"content": "# soft_dark_stylesheet.py\n\nsoft_dark_stylesheet = \"\"\"\nQWidget {\n background-color: #2F2F2F;\n color: #E0E0E0;\n font-family: Arial, sans-serif;\n}\n\nQMainWindow {\n background-color: #2A2A2A;\n}\n\nQPushButton {\n background-color: #4A4A4A;\n border: 1px solid #5E5E5E;\n padding: 5px 10px;\n border-radius: 3px;\n color: #E0E0E0;\n}\n\nQPushButton:hover {\n background-color: #545454;\n}\n\nQPushButton:pressed {\n background-color: #404040;\n}\n\nQPushButton:checked {\n background-color: #606060;\n border: 2px solid #808080;\n color: #FFFFFF;\n}\n\nQListWidget, QTreeWidget {\n background-color: #3A3A3A;\n border: 1px solid #4A4A4A;\n border-radius: 3px;\n color: #E0E0E0;\n}\n\nQListWidget::item, QTreeWidget::item {\n color: #E0E0E0; \n}\n\nQListWidget::item:selected, QTreeWidget::item:selected {\n background-color: #4A4A4A;\n color: #FFFFFF; /* Make selected items a bit brighter */\n}\n\nQLabel {\n color: #E0E0E0;\n}\n\nQLabel.section-header {\n font-weight: bold;\n font-size: 14px;\n padding: 5px 0;\n color: #FFFFFF; /* Bright white color for better visibility in dark mode */\n}\n\nQLineEdit, QTextEdit, QPlainTextEdit {\n background-color: #3A3A3A;\n border: 1px solid #4A4A4A;\n color: #E0E0E0;\n padding: 2px;\n border-radius: 3px;\n}\n\nQSlider::groove:horizontal {\n background: #4A4A4A;\n height: 8px;\n border-radius: 4px;\n}\n\nQSlider::handle:horizontal {\n background: #6A6A6A;\n width: 18px;\n margin-top: -5px;\n margin-bottom: -5px;\n border-radius: 9px;\n}\n\nQSlider::handle:horizontal:hover {\n background: #7A7A7A;\n}\n\nQScrollBar:vertical, QScrollBar:horizontal {\n background-color: #3A3A3A;\n width: 12px;\n height: 12px;\n}\n\nQScrollBar::handle:vertical, QScrollBar::handle:horizontal {\n background-color: #5A5A5A;\n border-radius: 6px;\n min-height: 20px;\n}\n\nQScrollBar::handle:vertical:hover, QScrollBar::handle:horizontal:hover {\n background-color: #6A6A6A;\n}\n\nQScrollBar::add-line, QScrollBar::sub-line {\n background: none;\n}\n\nQMenuBar {\n background-color: #2F2F2F;\n}\n\nQMenuBar::item {\n padding: 5px 10px;\n background-color: transparent;\n}\n\nQMenuBar::item:selected {\n background-color: #3A3A3A;\n}\n\nQMenu {\n background-color: #2F2F2F;\n border: 1px solid #3A3A3A;\n}\n\nQMenu::item {\n padding: 5px 20px 5px 20px;\n}\n\nQMenu::item:selected {\n background-color: #3A3A3A;\n}\n\nQToolTip {\n background-color: #2F2F2F;\n color: #E0E0E0;\n border: 1px solid #3A3A3A;\n}\n\nQStatusBar {\n background-color: #2A2A2A;\n color: #B0B0B0;\n}\n\nQListWidget::item {\n color: none;\n}\n\"\"\""
},
{
"path": "src/digitalsreeni_image_annotator/stack_interpolator.py",
"content": "import os\nimport numpy as np\nfrom PyQt5.QtWidgets import (QDialog, QVBoxLayout, QHBoxLayout, QPushButton, QFileDialog, \n QLabel, QComboBox, QMessageBox, QProgressDialog, QRadioButton,\n QButtonGroup, QGroupBox, QDoubleSpinBox, QApplication)\nfrom PyQt5.QtCore import Qt\nfrom scipy.interpolate import RegularGridInterpolator\nfrom skimage import io\nimport tifffile\n\nclass StackInterpolator(QDialog):\n def __init__(self, parent=None):\n super().__init__(parent)\n self.setWindowTitle(\"Stack Interpolator\")\n self.setGeometry(100, 100, 600, 400)\n self.setWindowFlags(self.windowFlags() | Qt.Window)\n self.setWindowModality(Qt.ApplicationModal) # Added window modality\n \n # Initialize variables\n self.input_path = \"\"\n self.output_directory = \"\"\n \n self.initUI()\n\n def initUI(self):\n layout = QVBoxLayout()\n layout.setSpacing(10) # Add consistent spacing\n\n # Input selection\n input_group = QGroupBox(\"Input Selection\")\n input_layout = QVBoxLayout()\n \n # Radio buttons for input type\n self.dir_radio = QRadioButton(\"Directory of Image Files\")\n self.stack_radio = QRadioButton(\"TIFF Stack\")\n \n input_group_buttons = QButtonGroup(self)\n input_group_buttons.addButton(self.dir_radio)\n input_group_buttons.addButton(self.stack_radio)\n \n input_layout.addWidget(self.dir_radio)\n input_layout.addWidget(self.stack_radio)\n self.dir_radio.setChecked(True)\n \n input_group.setLayout(input_layout)\n layout.addWidget(input_group)\n\n # Interpolation method\n method_group = QGroupBox(\"Interpolation Settings\")\n method_layout = QVBoxLayout()\n \n method_combo_layout = QHBoxLayout()\n method_combo_layout.addWidget(QLabel(\"Method:\"))\n self.method_combo = QComboBox()\n self.method_combo.addItems([\n \"linear\",\n \"nearest\",\n \"slinear\",\n \"cubic\",\n \"quintic\",\n \"pchip\"\n ])\n method_combo_layout.addWidget(self.method_combo)\n method_layout.addLayout(method_combo_layout)\n \n method_group.setLayout(method_layout)\n layout.addWidget(method_group)\n\n # Original dimensions group\n orig_group = QGroupBox(\"Original Dimensions\")\n orig_layout = QVBoxLayout()\n \n orig_xy_layout = QHBoxLayout()\n orig_xy_layout.addWidget(QLabel(\"XY Pixel Size:\"))\n self.orig_xy_size = QDoubleSpinBox()\n self.orig_xy_size.setRange(0.001, 1000.0)\n self.orig_xy_size.setValue(1.0)\n self.orig_xy_size.setDecimals(3)\n orig_xy_layout.addWidget(self.orig_xy_size)\n \n orig_z_layout = QHBoxLayout()\n orig_z_layout.addWidget(QLabel(\"Z Spacing:\"))\n self.orig_z_size = QDoubleSpinBox()\n self.orig_z_size.setRange(0.001, 1000.0)\n self.orig_z_size.setValue(1.0)\n self.orig_z_size.setDecimals(3)\n orig_z_layout.addWidget(self.orig_z_size)\n \n orig_layout.addLayout(orig_xy_layout)\n orig_layout.addLayout(orig_z_layout)\n orig_group.setLayout(orig_layout)\n layout.addWidget(orig_group)\n\n # New dimensions group\n new_group = QGroupBox(\"New Dimensions\")\n new_layout = QVBoxLayout()\n \n new_xy_layout = QHBoxLayout()\n new_xy_layout.addWidget(QLabel(\"XY Pixel Size:\"))\n self.new_xy_size = QDoubleSpinBox()\n self.new_xy_size.setRange(0.001, 1000.0)\n self.new_xy_size.setValue(1.0)\n self.new_xy_size.setDecimals(3)\n new_xy_layout.addWidget(self.new_xy_size)\n \n new_z_layout = QHBoxLayout()\n new_z_layout.addWidget(QLabel(\"Z Spacing:\"))\n self.new_z_size = QDoubleSpinBox()\n self.new_z_size.setRange(0.001, 1000.0)\n self.new_z_size.setValue(1.0)\n self.new_z_size.setDecimals(3)\n new_z_layout.addWidget(self.new_z_size)\n \n new_layout.addLayout(new_xy_layout)\n new_layout.addLayout(new_z_layout)\n new_group.setLayout(new_layout)\n layout.addWidget(new_group)\n\n # Units selector\n unit_group = QGroupBox(\"Unit Settings\")\n unit_layout = QHBoxLayout()\n unit_layout.addWidget(QLabel(\"Unit:\"))\n self.size_unit = QComboBox()\n self.size_unit.addItems([\"nm\", \"µm\", \"mm\"])\n self.size_unit.setCurrentText(\"µm\")\n unit_layout.addWidget(self.size_unit)\n unit_group.setLayout(unit_layout)\n layout.addWidget(unit_group)\n\n # Input/Output buttons\n button_group = QGroupBox(\"File Selection\")\n button_layout = QVBoxLayout()\n \n # Input selection\n input_file_layout = QHBoxLayout()\n self.input_label = QLabel(\"No input selected\")\n self.select_input_btn = QPushButton(\"Select Input\")\n self.select_input_btn.clicked.connect(self.select_input)\n input_file_layout.addWidget(self.select_input_btn)\n input_file_layout.addWidget(self.input_label)\n button_layout.addLayout(input_file_layout)\n \n # Output selection\n output_file_layout = QHBoxLayout()\n self.output_label = QLabel(\"No output directory selected\")\n self.select_output_btn = QPushButton(\"Select Output Directory\")\n self.select_output_btn.clicked.connect(self.select_output)\n output_file_layout.addWidget(self.select_output_btn)\n output_file_layout.addWidget(self.output_label)\n button_layout.addLayout(output_file_layout)\n \n button_group.setLayout(button_layout)\n layout.addWidget(button_group)\n\n # Interpolate button\n self.interpolate_btn = QPushButton(\"Interpolate\")\n self.interpolate_btn.clicked.connect(self.interpolate_stack)\n layout.addWidget(self.interpolate_btn)\n\n self.setLayout(layout)\n\n def select_input(self):\n try:\n if self.dir_radio.isChecked():\n path = QFileDialog.getExistingDirectory(\n self,\n \"Select Directory with Images\",\n \"\",\n QFileDialog.ShowDirsOnly | QFileDialog.DontResolveSymlinks\n )\n else:\n path, _ = QFileDialog.getOpenFileName(\n self,\n \"Select TIFF Stack\",\n \"\",\n \"TIFF Files (*.tif *.tiff)\",\n options=QFileDialog.Options()\n )\n \n if path:\n self.input_path = path\n self.input_label.setText(f\"Selected: {os.path.basename(path)}\")\n self.input_label.setToolTip(path)\n QApplication.processEvents()\n \n except Exception as e:\n QMessageBox.critical(self, \"Error\", f\"Error selecting input: {str(e)}\")\n\n def select_output(self):\n try:\n directory = QFileDialog.getExistingDirectory(\n self,\n \"Select Output Directory\",\n \"\",\n QFileDialog.ShowDirsOnly | QFileDialog.DontResolveSymlinks\n )\n \n if directory:\n self.output_directory = directory\n self.output_label.setText(f\"Selected: {os.path.basename(directory)}\")\n self.output_label.setToolTip(directory)\n QApplication.processEvents()\n \n except Exception as e:\n QMessageBox.critical(self, \"Error\", f\"Error selecting output directory: {str(e)}\")\n\n def load_images(self):\n try:\n progress = QProgressDialog(\"Loading images...\", \"Cancel\", 0, 100, self)\n progress.setWindowModality(Qt.WindowModal)\n progress.show()\n QApplication.processEvents()\n \n if self.stack_radio.isChecked():\n progress.setLabelText(\"Loading TIFF stack...\")\n progress.setValue(20)\n QApplication.processEvents()\n \n # Load stack preserving original dtype\n stack = io.imread(self.input_path)\n print(f\"Loaded stack dtype: {stack.dtype}\")\n print(f\"Value range: [{stack.min()}, {stack.max()}]\")\n \n progress.setValue(90)\n QApplication.processEvents()\n return stack\n else:\n valid_extensions = ('.png', '.jpg', '.jpeg', '.bmp', '.tif', '.tiff')\n files = sorted([f for f in os.listdir(self.input_path) \n if f.lower().endswith(valid_extensions)])\n \n if not files:\n raise ValueError(\"No valid image files found in directory\")\n \n progress.setMaximum(len(files))\n \n # Load first image to get dimensions and dtype\n first_img = io.imread(os.path.join(self.input_path, files[0]))\n stack = np.zeros((len(files), *first_img.shape), dtype=first_img.dtype)\n stack[0] = first_img\n \n print(f\"Created stack with dtype: {stack.dtype}\")\n print(f\"First image range: [{first_img.min()}, {first_img.max()}]\")\n \n # Load remaining images\n for i, fname in enumerate(files[1:], 1):\n progress.setValue(i)\n progress.setLabelText(f\"Loading image {i+1}/{len(files)}\")\n QApplication.processEvents()\n \n if progress.wasCanceled():\n raise InterruptedError(\"Loading cancelled by user\")\n \n img = io.imread(os.path.join(self.input_path, fname))\n if img.shape != first_img.shape:\n raise ValueError(f\"Image {fname} has different dimensions from the first image\")\n if img.dtype != first_img.dtype:\n raise ValueError(f\"Image {fname} has different bit depth from the first image\")\n stack[i] = img\n \n return stack\n \n except Exception as e:\n raise ValueError(f\"Error loading images: {str(e)}\")\n finally:\n progress.close()\n QApplication.processEvents()\n \n def interpolate_stack(self):\n if not self.input_path or not self.output_directory:\n QMessageBox.warning(self, \"Missing Paths\", \"Please select both input and output paths\")\n return\n \n try:\n # Create progress dialog\n progress = QProgressDialog(\"Processing...\", \"Cancel\", 0, 100, self)\n progress.setWindowModality(Qt.WindowModal)\n progress.setWindowTitle(\"Interpolation Progress\")\n progress.setMinimumDuration(0)\n progress.setMinimumWidth(400)\n progress.show()\n QApplication.processEvents()\n \n # Load images\n progress.setLabelText(\"Loading images...\")\n progress.setValue(10)\n QApplication.processEvents()\n \n input_stack = self.load_images()\n original_dtype = input_stack.dtype\n type_range = np.iinfo(original_dtype) if np.issubdtype(original_dtype, np.integer) else None\n \n print(f\"Original data type: {original_dtype}\")\n print(f\"Original shape: {input_stack.shape}\")\n print(f\"Original range: {input_stack.min()} - {input_stack.max()}\")\n \n # Normalize input data to float64 for interpolation\n input_stack_normalized = input_stack.astype(np.float64)\n if type_range is not None:\n input_stack_normalized = input_stack_normalized / type_range.max\n \n progress.setLabelText(\"Calculating dimensions...\")\n progress.setValue(20)\n QApplication.processEvents()\n \n # Calculate dimensions and coordinates\n z_old = np.arange(input_stack.shape[0]) * self.orig_z_size.value()\n y_old = np.arange(input_stack.shape[1]) * self.orig_xy_size.value()\n x_old = np.arange(input_stack.shape[2]) * self.orig_xy_size.value()\n \n z_new = np.arange(z_old[0], z_old[-1] + self.new_z_size.value(), self.new_z_size.value())\n y_new = np.arange(0, input_stack.shape[1] * self.orig_xy_size.value(), self.new_xy_size.value())\n x_new = np.arange(0, input_stack.shape[2] * self.orig_xy_size.value(), self.new_xy_size.value())\n \n y_new = y_new[y_new < y_old[-1] + self.new_xy_size.value()]\n x_new = x_new[x_new < x_old[-1] + self.new_xy_size.value()]\n \n new_shape = (len(z_new), len(y_new), len(x_new))\n print(f\"New dimensions will be: {new_shape}\")\n \n # Initialize output array\n interpolated_data = np.zeros(new_shape, dtype=np.float64)\n \n method = self.method_combo.currentText()\n \n # For higher-order methods, use a hybrid approach\n if method in ['cubic', 'quintic', 'pchip']:\n progress.setLabelText(\"Using hybrid interpolation approach...\")\n progress.setValue(30)\n QApplication.processEvents()\n \n from scipy.interpolate import interp1d\n \n # Process each XY point\n total_points = input_stack.shape[1] * input_stack.shape[2]\n points_processed = 0\n \n temp_stack = np.zeros((len(z_new), input_stack.shape[1], input_stack.shape[2]), dtype=np.float64)\n \n for y in range(input_stack.shape[1]):\n for x in range(input_stack.shape[2]):\n if progress.wasCanceled():\n return\n \n points_processed += 1\n if points_processed % 1000 == 0:\n progress_val = 30 + (points_processed / total_points * 30)\n progress.setValue(int(progress_val))\n progress.setLabelText(f\"Interpolating Z dimension: {points_processed}/{total_points} points\")\n QApplication.processEvents()\n \n z_profile = input_stack_normalized[:, y, x]\n f = interp1d(z_old, z_profile, kind=method, bounds_error=False, fill_value='extrapolate')\n temp_stack[:, y, x] = f(z_new)\n \n progress.setLabelText(\"Interpolating XY planes...\")\n progress.setValue(60)\n QApplication.processEvents()\n \n for z in range(len(z_new)):\n if progress.wasCanceled():\n return\n \n progress.setValue(60 + int((z / len(z_new)) * 30))\n progress.setLabelText(f\"Processing XY plane {z+1}/{len(z_new)}\")\n QApplication.processEvents()\n \n interpolator = RegularGridInterpolator(\n (y_old, x_old),\n temp_stack[z],\n method='linear',\n bounds_error=False,\n fill_value=0\n )\n \n yy, xx = np.meshgrid(y_new, x_new, indexing='ij')\n pts = np.stack([yy.ravel(), xx.ravel()], axis=-1)\n \n interpolated_data[z] = interpolator(pts).reshape(len(y_new), len(x_new))\n \n del temp_stack\n \n else: # For linear and nearest neighbor\n progress.setLabelText(\"Creating interpolator...\")\n progress.setValue(30)\n QApplication.processEvents()\n \n interpolator = RegularGridInterpolator(\n (z_old, y_old, x_old),\n input_stack_normalized,\n method=method,\n bounds_error=False,\n fill_value=0\n )\n \n slices_per_batch = max(1, len(z_new) // 20)\n total_batches = (len(z_new) + slices_per_batch - 1) // slices_per_batch\n \n for batch_idx in range(total_batches):\n if progress.wasCanceled():\n return\n \n start_idx = batch_idx * slices_per_batch\n end_idx = min((batch_idx + 1) * slices_per_batch, len(z_new))\n \n progress.setLabelText(f\"Interpolating batch {batch_idx + 1}/{total_batches}\")\n progress_value = int(40 + (batch_idx/total_batches)*40)\n progress.setValue(progress_value)\n QApplication.processEvents()\n \n zz, yy, xx = np.meshgrid(\n z_new[start_idx:end_idx],\n y_new,\n x_new,\n indexing='ij'\n )\n \n pts = np.stack([zz.ravel(), yy.ravel(), xx.ravel()], axis=-1)\n \n interpolated_data[start_idx:end_idx] = interpolator(pts).reshape(\n end_idx - start_idx,\n len(y_new),\n len(x_new)\n )\n \n # Convert back to original dtype\n progress.setLabelText(\"Converting to original bit depth...\")\n progress.setValue(90)\n QApplication.processEvents()\n \n if np.issubdtype(original_dtype, np.integer):\n # Scale back to original range\n interpolated_data = np.clip(interpolated_data, 0, 1)\n interpolated_data = (interpolated_data * type_range.max).astype(original_dtype)\n else:\n interpolated_data = interpolated_data.astype(original_dtype)\n \n print(f\"Final dtype: {interpolated_data.dtype}\")\n print(f\"Final range: [{interpolated_data.min()}, {interpolated_data.max()}]\")\n \n # Save output\n progress.setLabelText(\"Saving interpolated stack...\")\n progress.setValue(95)\n QApplication.processEvents()\n \n if self.stack_radio.isChecked():\n output_name = os.path.splitext(os.path.basename(self.input_path))[0]\n else:\n output_name = \"interpolated_stack\"\n \n output_path = os.path.join(self.output_directory, f\"{output_name}_interpolated.tif\")\n \n # Convert sizes to micrometers for metadata\n unit = self.size_unit.currentText()\n xy_size = self.new_xy_size.value()\n z_size = self.new_z_size.value()\n \n if unit == \"nm\":\n xy_size /= 1000\n z_size /= 1000\n elif unit == \"mm\":\n xy_size *= 1000\n z_size *= 1000\n \n # Save with metadata\n tifffile.imwrite(\n output_path,\n interpolated_data,\n imagej=True,\n metadata={\n 'axes': 'ZYX',\n 'spacing': z_size,\n 'unit': 'um',\n 'finterval': xy_size\n },\n resolution=(1.0/xy_size, 1.0/xy_size)\n )\n \n progress.setValue(100)\n QApplication.processEvents()\n \n QMessageBox.information(\n self,\n \"Success\",\n f\"Interpolation completed successfully!\\n\"\n f\"Output saved to:\\n{output_path}\\n\"\n f\"New dimensions: {interpolated_data.shape}\\n\"\n f\"Bit depth: {interpolated_data.dtype}\\n\"\n f\"XY Pixel size: {self.new_xy_size.value()} {unit}\\n\"\n f\"Z Spacing: {self.new_z_size.value()} {unit}\"\n )\n \n except Exception as e:\n QMessageBox.critical(self, \"Error\", str(e))\n print(f\"Error occurred: {str(e)}\")\n import traceback\n traceback.print_exc()\n finally:\n progress.close()\n QApplication.processEvents()\n\n\n def show_centered(self, parent):\n parent_geo = parent.geometry()\n self.move(parent_geo.center() - self.rect().center())\n self.show()\n QApplication.processEvents() # Ensure UI updates\n \n \n# Helper function to create the dialog\ndef show_stack_interpolator(parent):\n dialog = StackInterpolator(parent)\n dialog.show_centered(parent)\n return dialog"
},
{
"path": "src/digitalsreeni_image_annotator/stack_to_slices.py",
"content": "import os\nimport numpy as np\nfrom PyQt5.QtWidgets import (QDialog, QVBoxLayout, QHBoxLayout, QPushButton, \n QFileDialog, QLabel, QMessageBox, QComboBox, QGridLayout, QWidget,\n QProgressDialog, QApplication)\nfrom PyQt5.QtCore import Qt\nfrom tifffile import TiffFile\nfrom czifile import CziFile\nfrom PIL import Image\n\nclass DimensionDialog(QDialog):\n def __init__(self, shape, file_name, parent=None):\n super().__init__(parent)\n self.setWindowTitle(\"Assign Dimensions\")\n self.shape = shape\n self.initUI(file_name)\n\n def initUI(self, file_name):\n layout = QVBoxLayout()\n \n file_name_label = QLabel(f\"File: {file_name}\")\n file_name_label.setWordWrap(True)\n layout.addWidget(file_name_label)\n \n dim_widget = QWidget()\n dim_layout = QGridLayout(dim_widget)\n self.combos = []\n dimensions = ['T', 'Z', 'C', 'S', 'H', 'W']\n for i, dim in enumerate(self.shape):\n dim_layout.addWidget(QLabel(f\"Dimension {i} (size {dim}):\"), i, 0)\n combo = QComboBox()\n combo.addItems(dimensions)\n dim_layout.addWidget(combo, i, 1)\n self.combos.append(combo)\n layout.addWidget(dim_widget)\n \n self.button = QPushButton(\"OK\")\n self.button.clicked.connect(self.accept)\n layout.addWidget(self.button)\n \n self.setLayout(layout)\n\n def get_dimensions(self):\n return [combo.currentText() for combo in self.combos]\n\n\nclass StackToSlicesDialog(QDialog):\n def __init__(self, parent=None):\n super().__init__(parent)\n self.setWindowTitle(\"Stack to Slices\")\n self.setGeometry(100, 100, 400, 200)\n self.setWindowFlags(self.windowFlags() | Qt.Window)\n self.setWindowModality(Qt.ApplicationModal)\n self.dimensions = None\n self.initUI()\n\n def initUI(self):\n layout = QVBoxLayout()\n\n self.file_label = QLabel(\"No file selected\")\n layout.addWidget(self.file_label)\n\n select_button = QPushButton(\"Select Stack File\")\n select_button.clicked.connect(self.select_file)\n layout.addWidget(select_button)\n\n self.convert_button = QPushButton(\"Convert to Slices\")\n self.convert_button.clicked.connect(self.convert_to_slices)\n self.convert_button.setEnabled(False)\n layout.addWidget(self.convert_button)\n\n self.setLayout(layout)\n\n def select_file(self):\n self.file_name, _ = QFileDialog.getOpenFileName(self, \"Select Stack File\", \"\", \"Image Files (*.tif *.tiff *.czi)\")\n if self.file_name:\n self.file_label.setText(f\"Selected file: {os.path.basename(self.file_name)}\")\n QApplication.processEvents()\n self.process_file()\n\n def process_file(self):\n if self.file_name.lower().endswith(('.tif', '.tiff')):\n self.process_tiff()\n elif self.file_name.lower().endswith('.czi'):\n self.process_czi()\n\n def process_tiff(self):\n with TiffFile(self.file_name) as tif:\n image_array = tif.asarray()\n \n self.get_dimensions(image_array.shape)\n\n def process_czi(self):\n with CziFile(self.file_name) as czi:\n image_array = czi.asarray()\n \n self.get_dimensions(image_array.shape)\n\n def get_dimensions(self, shape):\n dialog = DimensionDialog(shape, os.path.basename(self.file_name), self)\n dialog.setWindowModality(Qt.ApplicationModal)\n if dialog.exec_():\n self.dimensions = dialog.get_dimensions()\n self.convert_button.setEnabled(True)\n else:\n self.dimensions = None\n self.convert_button.setEnabled(False)\n QApplication.processEvents()\n\n def convert_to_slices(self):\n if not hasattr(self, 'file_name') or not self.dimensions:\n QMessageBox.warning(self, \"Invalid Input\", \"Please select a file and assign dimensions first.\")\n return\n\n output_dir = QFileDialog.getExistingDirectory(self, \"Select Output Directory\")\n if not output_dir:\n return\n\n if self.file_name.lower().endswith(('.tif', '.tiff')):\n with TiffFile(self.file_name) as tif:\n image_array = tif.asarray()\n elif self.file_name.lower().endswith('.czi'):\n with CziFile(self.file_name) as czi:\n image_array = czi.asarray()\n\n self.save_slices(image_array, output_dir)\n\n def save_slices(self, image_array, output_dir):\n base_name = os.path.splitext(os.path.basename(self.file_name))[0]\n \n slice_indices = [i for i, dim in enumerate(self.dimensions) if dim not in ['H', 'W']]\n\n total_slices = np.prod([image_array.shape[i] for i in slice_indices])\n \n progress = QProgressDialog(\"Saving slices...\", \"Cancel\", 0, total_slices, self)\n progress.setWindowModality(Qt.WindowModal)\n progress.setWindowTitle(\"Progress\")\n progress.setMinimumDuration(0)\n progress.setValue(0)\n progress.show()\n\n try:\n for idx, _ in enumerate(np.ndindex(tuple(image_array.shape[i] for i in slice_indices))):\n if progress.wasCanceled():\n break\n\n full_idx = [slice(None)] * len(self.dimensions)\n for i, val in zip(slice_indices, _):\n full_idx[i] = val\n \n slice_array = image_array[tuple(full_idx)]\n \n if slice_array.ndim > 2:\n slice_array = slice_array.squeeze()\n \n if slice_array.dtype == np.uint16:\n mode = 'I;16'\n elif slice_array.dtype == np.uint8:\n mode = 'L'\n else:\n slice_array = ((slice_array - slice_array.min()) / (slice_array.max() - slice_array.min()) * 65535).astype(np.uint16)\n mode = 'I;16'\n\n slice_name = f\"{base_name}_{'_'.join([f'{self.dimensions[i]}{val+1}' for i, val in zip(slice_indices, _)])}.png\"\n img = Image.fromarray(slice_array, mode=mode)\n img.save(os.path.join(output_dir, slice_name))\n\n progress.setValue(idx + 1)\n QApplication.processEvents()\n\n if progress.wasCanceled():\n QMessageBox.warning(self, \"Conversion Interrupted\", \"The conversion process was interrupted.\")\n else:\n QMessageBox.information(self, \"Conversion Complete\", f\"All slices have been saved to {output_dir}\")\n \n finally:\n progress.close()\n\n def show_centered(self, parent):\n parent_geo = parent.geometry()\n self.move(parent_geo.center() - self.rect().center())\n self.show()\n\ndef show_stack_to_slices(parent):\n dialog = StackToSlicesDialog(parent)\n dialog.show_centered(parent)\n return dialog"
},
{
"path": "src/digitalsreeni_image_annotator/utils.py",
"content": "\"\"\"\nUtility functions for the Image Annotator application.\n\nThis module contains helper functions used across the application.\n\n@DigitalSreeni\nDr. Sreenivas Bhattiprolu\n\"\"\"\n\nimport numpy as np\n\ndef calculate_area(annotation):\n if \"segmentation\" in annotation and annotation[\"segmentation\"] is not None:\n # Polygon area\n x, y = annotation[\"segmentation\"][0::2], annotation[\"segmentation\"][1::2]\n return 0.5 * abs(sum(x[i] * y[i+1] - x[i+1] * y[i] for i in range(-1, len(x)-1)))\n elif \"bbox\" in annotation:\n # Rectangle area\n x, y, w, h = annotation[\"bbox\"]\n return w * h\n return 0\n\ndef calculate_bbox(segmentation):\n x_coordinates, y_coordinates = segmentation[0::2], segmentation[1::2]\n x_min, y_min = min(x_coordinates), min(y_coordinates)\n x_max, y_max = max(x_coordinates), max(y_coordinates)\n width, height = x_max - x_min, y_max - y_min\n return [x_min, y_min, width, height]\n\ndef normalize_image(image_array):\n \"\"\"Normalize image array to 8-bit range.\"\"\"\n if image_array.dtype != np.uint8:\n image_array = ((image_array - image_array.min()) / (image_array.max() - image_array.min()) * 255).astype(np.uint8)\n return image_array\n\n"
},
{
"path": "src/digitalsreeni_image_annotator/yolo_trainer.py",
"content": "import os\nfrom ultralytics import YOLO\nfrom PyQt5.QtWidgets import QFileDialog, QMessageBox\nfrom PyQt5.QtWidgets import (QDialog, QVBoxLayout, QHBoxLayout, QPushButton, \n QLineEdit, QLabel, QFileDialog, QDialogButtonBox)\nimport yaml\nimport numpy as np\nfrom pathlib import Path\nfrom .export_formats import export_yolo_v5plus\n\n\nfrom collections import deque\n\n\nfrom PyQt5.QtWidgets import QDialog, QVBoxLayout, QTextEdit, QPushButton\nfrom PyQt5.QtCore import Qt, pyqtSignal, QObject\n\nclass TrainingInfoDialog(QDialog):\n stop_signal = pyqtSignal()\n\n def __init__(self, parent=None):\n super().__init__(parent)\n self.setWindowTitle(\"Training Progress\")\n self.setModal(False)\n self.layout = QVBoxLayout(self)\n\n self.info_text = QTextEdit(self)\n self.info_text.setReadOnly(True)\n self.layout.addWidget(self.info_text)\n\n self.stop_button = QPushButton(\"Stop Training\", self)\n self.stop_button.clicked.connect(self.stop_training)\n self.layout.addWidget(self.stop_button)\n\n self.close_button = QPushButton(\"Close\", self)\n self.close_button.clicked.connect(self.hide)\n self.layout.addWidget(self.close_button)\n\n self.setMinimumSize(400, 300)\n\n def update_info(self, text):\n self.info_text.append(text)\n self.info_text.verticalScrollBar().setValue(self.info_text.verticalScrollBar().maximum())\n\n def stop_training(self):\n self.stop_signal.emit()\n self.stop_button.setEnabled(False)\n self.stop_button.setText(\"Stopping...\")\n\n def closeEvent(self, event):\n event.ignore()\n self.hide()\n\nclass LoadPredictionModelDialog(QDialog):\n def __init__(self, parent=None):\n super().__init__(parent)\n self.setWindowTitle(\"Load Prediction Model and YAML\")\n self.model_path = \"\"\n self.yaml_path = \"\"\n\n layout = QVBoxLayout(self)\n\n # Model file selection\n model_layout = QHBoxLayout()\n self.model_edit = QLineEdit()\n model_button = QPushButton(\"Browse\")\n model_button.clicked.connect(self.browse_model)\n model_layout.addWidget(QLabel(\"Model File:\"))\n model_layout.addWidget(self.model_edit)\n model_layout.addWidget(model_button)\n layout.addLayout(model_layout)\n\n # YAML file selection\n yaml_layout = QHBoxLayout()\n self.yaml_edit = QLineEdit()\n yaml_button = QPushButton(\"Browse\")\n yaml_button.clicked.connect(self.browse_yaml)\n yaml_layout.addWidget(QLabel(\"YAML File:\"))\n yaml_layout.addWidget(self.yaml_edit)\n yaml_layout.addWidget(yaml_button)\n layout.addLayout(yaml_layout)\n\n # OK and Cancel buttons\n self.button_box = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel)\n self.button_box.accepted.connect(self.accept)\n self.button_box.rejected.connect(self.reject)\n layout.addWidget(self.button_box)\n\n def browse_model(self):\n file_name, _ = QFileDialog.getOpenFileName(self, \"Select YOLO Model\", \"\", \"YOLO Model (*.pt)\")\n if file_name:\n self.model_path = file_name\n self.model_edit.setText(file_name)\n\n def browse_yaml(self):\n file_name, _ = QFileDialog.getOpenFileName(self, \"Select YAML File\", \"\", \"YAML Files (*.yaml *.yml)\")\n if file_name:\n self.yaml_path = file_name\n self.yaml_edit.setText(file_name)\n \nclass YOLOTrainer(QObject):\n progress_signal = pyqtSignal(str)\n\n def __init__(self, project_dir, main_window):\n super().__init__()\n self.project_dir = project_dir\n self.main_window = main_window\n self.model = None\n self.dataset_path = os.path.join(project_dir, \"yolo_dataset\")\n self.model_path = os.path.join(project_dir, \"yolo_model\")\n self.yaml_path = None\n self.yaml_data = None\n self.epoch_info = deque(maxlen=10)\n self.progress_callback = None\n self.total_epochs = None\n self.conf_threshold = 0.25\n self.stop_training = False\n self.class_names = None\n\n def load_model(self, model_path=None):\n if model_path is None:\n model_path, _ = QFileDialog.getOpenFileName(self.main_window, \"Select YOLO Model\", \"\", \"YOLO Model (*.pt)\")\n if model_path:\n try:\n self.model = YOLO(model_path)\n return True\n except Exception as e:\n QMessageBox.critical(self.main_window, \"Error Loading Model\", f\"Could not load the model. Error: {str(e)}\")\n return False\n\n def prepare_dataset(self):\n output_dir, yaml_path = export_yolo_v5plus(\n self.main_window.all_annotations,\n self.main_window.class_mapping,\n self.main_window.image_paths,\n self.main_window.slices,\n self.main_window.image_slices,\n self.dataset_path\n )\n \n yaml_path = Path(yaml_path)\n with yaml_path.open('r') as f:\n yaml_content = yaml.safe_load(f)\n \n # Update paths for new YOLO v5+ structure\n yaml_content['train'] = 'images/train' # Changed from train/images\n yaml_content['val'] = 'images/val' # Changed from train/images\n yaml_content['test'] = '../test/images'\n \n with yaml_path.open('w') as f:\n yaml.dump(yaml_content, f, default_flow_style=False)\n \n self.yaml_path = str(yaml_path)\n return self.yaml_path\n\n def load_yaml(self, yaml_path=None):\n if yaml_path is None:\n yaml_path, _ = QFileDialog.getOpenFileName(self.main_window, \"Select YOLO Dataset YAML\", \"\", \"YAML Files (*.yaml *.yml)\")\n if yaml_path and os.path.exists(yaml_path):\n with open(yaml_path, 'r') as f:\n try:\n yaml_data = yaml.safe_load(f)\n print(f\"Loaded YAML contents: {yaml_data}\")\n \n # Ensure paths are relative\n for key in ['train', 'val', 'test']:\n if key in yaml_data and os.path.isabs(yaml_data[key]):\n yaml_data[key] = os.path.relpath(yaml_data[key], start=os.path.dirname(yaml_path))\n \n print(f\"Updated YAML contents: {yaml_data}\")\n \n # Save the updated YAML data\n self.yaml_data = yaml_data\n self.yaml_path = yaml_path\n \n # Write the updated YAML back to the file\n with open(yaml_path, 'w') as f:\n yaml.dump(yaml_data, f, default_flow_style=False)\n \n return True\n except yaml.YAMLError as e:\n QMessageBox.critical(self.main_window, \"Error Loading YAML\", f\"Invalid YAML file. Error: {str(e)}\")\n return False\n\n def on_train_epoch_end(self, trainer):\n epoch = trainer.epoch + 1 # Add 1 to start from 1 instead of 0\n total_epochs = trainer.epochs\n loss = trainer.loss.item()\n progress_text = f\"Epoch {epoch}/{total_epochs}, Loss: {loss:.4f}\"\n \n # Only emit the signal, don't call the callback directly\n self.progress_signal.emit(progress_text)\n \n if self.stop_training:\n trainer.model.stop = True\n self.stop_training = False\n return False\n return True\n \n def train_model(self, epochs=100, imgsz=640):\n if self.model is None:\n raise ValueError(\"No model loaded. Please load a model first.\")\n if self.yaml_path is None or not Path(self.yaml_path).exists():\n raise FileNotFoundError(\"Dataset YAML not found. Please prepare or load a dataset first.\")\n \n self.stop_training = False\n self.total_epochs = epochs\n self.epoch_info.clear()\n \n # Add the callback\n self.model.add_callback(\"on_train_epoch_end\", self.on_train_epoch_end)\n \n try:\n yaml_path = Path(self.yaml_path)\n yaml_dir = yaml_path.parent\n \n print(f\"Training with YAML: {yaml_path}\")\n print(f\"YAML directory: {yaml_dir}\")\n \n with yaml_path.open('r') as f:\n yaml_content = yaml.safe_load(f)\n print(f\"YAML content: {yaml_content}\")\n \n # For now, use train as val since we don't have separate validation set\n train_dir = str(yaml_dir / 'images' / 'train')\n \n # Update YAML content with correct paths\n yaml_content['train'] = train_dir\n yaml_content['val'] = train_dir # Use same directory for validation\n \n # Create the val directory structure if it doesn't exist\n val_img_dir = yaml_dir / 'images' / 'val'\n val_label_dir = yaml_dir / 'labels' / 'val'\n val_img_dir.mkdir(parents=True, exist_ok=True)\n val_label_dir.mkdir(parents=True, exist_ok=True)\n \n # Write updated YAML with adjusted paths\n temp_yaml_path = yaml_dir / 'temp_train.yaml'\n with temp_yaml_path.open('w') as f:\n yaml.dump(yaml_content, f, default_flow_style=False)\n \n print(f\"Training with updated YAML: {temp_yaml_path}\")\n print(f\"Updated YAML content: {yaml_content}\")\n \n results = self.model.train(data=str(temp_yaml_path), epochs=epochs, imgsz=imgsz)\n return results\n finally:\n # Clear the callback\n self.model.callbacks[\"on_train_epoch_end\"] = []\n # Remove temporary YAML file\n if 'temp_yaml_path' in locals():\n temp_yaml_path.unlink(missing_ok=True)\n \n def verify_dataset_structure(self):\n yaml_path = Path(self.yaml_path)\n yaml_dir = yaml_path.parent\n \n with yaml_path.open('r') as f:\n yaml_content = yaml.safe_load(f)\n \n # Use paths from YAML content\n train_images_dir = yaml_dir / yaml_content.get('train', 'images/train')\n val_images_dir = yaml_dir / yaml_content.get('val', 'images/val')\n train_labels_dir = yaml_dir / 'labels' / 'train' # Labels directory corresponds to images\n val_labels_dir = yaml_dir / 'labels' / 'val' # Labels directory corresponds to images\n \n # Check both train and val directories\n missing_dirs = []\n if not train_images_dir.exists():\n missing_dirs.append(f\"Training images directory: {train_images_dir}\")\n if not train_labels_dir.exists():\n missing_dirs.append(f\"Training labels directory: {train_labels_dir}\")\n if not val_images_dir.exists():\n missing_dirs.append(f\"Validation images directory: {val_images_dir}\")\n if not val_labels_dir.exists():\n missing_dirs.append(f\"Validation labels directory: {val_labels_dir}\")\n \n if missing_dirs:\n raise FileNotFoundError(f\"The following directories were not found:\\n\" + \"\\n\".join(missing_dirs))\n \n print(f\"Dataset structure verified:\")\n print(f\"Train images: {train_images_dir}\")\n print(f\"Train labels: {train_labels_dir}\")\n print(f\"Val images: {val_images_dir}\")\n print(f\"Val labels: {val_labels_dir}\")\n\n def check_ultralytics_settings(self):\n settings_path = Path.home() / \".config\" / \"Ultralytics\" / \"settings.yaml\"\n if settings_path.exists():\n with settings_path.open('r') as f:\n settings = yaml.safe_load(f)\n print(f\"Ultralytics settings: {settings}\")\n else:\n print(\"Ultralytics settings file not found.\")\n \n def stop_training_signal(self):\n self.stop_training = True\n self.progress_signal.emit(\"Stopping training...\")\n\n def set_progress_callback(self, callback):\n self.progress_callback = callback\n\n \n def stop_training_callback(self, trainer):\n if getattr(self, 'stop_training', False):\n trainer.model.stop = True\n self.stop_training = False\n \n\n \n def on_epoch_end(self, trainer):\n # Get current epoch\n epoch = trainer.epoch if hasattr(trainer, 'epoch') else trainer.current_epoch\n\n # Get total epochs\n total_epochs = self.total_epochs # Use the value we set in train_model\n\n # Get loss\n if hasattr(trainer, 'metrics') and 'train/box_loss' in trainer.metrics:\n loss = trainer.metrics['train/box_loss']\n elif hasattr(trainer, 'loss'):\n loss = trainer.loss\n else:\n loss = 0 # Default value if loss can't be found\n\n # Ensure loss is a number\n loss = float(loss)\n\n info = f\"Epoch {epoch}/{total_epochs}, Loss: {loss:.4f}\"\n self.epoch_info.append(info)\n \n display_text = f\"Current Progress:\\n\" + \"\\n\".join(self.epoch_info)\n if self.progress_callback:\n self.progress_callback(display_text)\n\n\n def save_model(self):\n if self.model is None:\n raise ValueError(\"No model to save. Please train a model first.\")\n save_path, _ = QFileDialog.getSaveFileName(self.main_window, \"Save YOLO Model\", \"\", \"YOLO Model (*.pt)\")\n if save_path:\n self.model.export(save_path)\n return True\n return False\n\n def load_prediction_model(self, model_path, yaml_path):\n try:\n self.model = YOLO(model_path)\n with open(yaml_path, 'r') as f:\n self.prediction_yaml = yaml.safe_load(f)\n \n if 'names' not in self.prediction_yaml:\n raise ValueError(\"The YAML file does not contain a 'names' section for class names.\")\n \n self.class_names = self.prediction_yaml['names']\n print(f\"Loaded class names: {self.class_names}\")\n \n # Verify that the number of classes in the YAML matches the model\n if len(self.class_names) != len(self.model.names):\n mismatch_message = (f\"Warning: Number of classes in YAML ({len(self.class_names)}) \"\n f\"does not match the model ({len(self.model.names)}). \"\n \"This may cause issues during prediction.\")\n print(mismatch_message)\n return True, mismatch_message\n \n return True, None\n except Exception as e:\n error_message = f\"Error loading model or YAML: {str(e)}\"\n print(error_message)\n return False, error_message\n \n def predict(self, input_data):\n if self.model is None:\n raise ValueError(\"No model loaded. Please load a model first.\")\n if isinstance(input_data, str):\n # It's a file path\n results = self.model(input_data, task='segment', conf=self.conf_threshold, save=False, show=False)\n elif isinstance(input_data, np.ndarray):\n # It's a numpy array\n results = self.model(input_data, task='segment', conf=self.conf_threshold, save=False, show=False)\n else:\n raise ValueError(\"Invalid input type. Expected file path or numpy array.\")\n \n # Get the input size used for prediction and the original image size\n input_size = results[0].orig_shape\n original_size = results[0].orig_img.shape[:2]\n return results, input_size, original_size\n\n def set_conf_threshold(self, conf):\n self.conf_threshold = conf"
}
]