Merge pull request #1 from electro199/dev

Add AutoLabelling and progress bar

.gitignore

@@ -6,8 +6,8 @@ test.python
 output/
 *.csv
 *.spec
-*.txt
 *.imtag
+label*.txt
 compile.py
 __old_code.py
 test.py

Readme.md

@@ -9,40 +9,48 @@ This Python-based Image Classification App allows you to label images for machin
 - **Folder Input:** Load images directly from a specified folder.
 - **CSV Input:** Use a CSV file to specify image paths and additional metadata.
 - **Progress Tracking:** Keep track of your progress with a progress file.
-- **Autosave** saving every changes into progress file automatically
+- **Autosave** saving every changes into progress file automatically.
+- **AutoLabelling** (Beta) Builitin support for auto lable prediction model running with transformers.
+
 ## Installation
 
 To install and set up the Image Classification App, follow these steps:
 
 1. **Clone the repository:**
 ``` bash
-   git clone https://github.com/electro199/imTagger.git
-   cd imTagger
-
+git clone https://github.com/electro199/imTagger.git
+cd imTagger
 ```
 
 
-install requirements
+Install requirements
 
+Python version 3.10 and 3.11 are recommended.
 
-
-```bash
 windows :
+```bash
 pip install -r requirements.txt
-
+```
 Unix :
+```bash
 pip3 install -r ./requirements.txt
+```
+
+## Auto Labelling
 
+This Feature is in beta to use in app set `AUTO_LABELER_ENABLED` to True in `main.py` and install dependency:
+
+```
+pip install transformers
 ```
 
+
 ## Usage :
 
 On windows click `run.bat` or in cmd `run.bat`
 
 On unix in terminal `python3 main.py`
 
 
-
-
 # Contributing
 Pull requests are more than welcome! If you are planning to contribute a large patch, please create an issue first to get any upfront questions or design decisions out of the way first.

auto_labeler.py

@@ -0,0 +1,51 @@
+import sys
+from PySide6.QtCore import QThread, Signal, Slot
+from PIL import Image
+import torch
+class Auto_Labbeler(QThread):
+    result_signal = Signal(str)  # Signal to send data back to the main thread
+    finished_signal = Signal()  # Signal to indicate that processing is finished
+
+    def __init__(self, parent=None):
+        super().__init__(parent)
+        self.model = None
+        self.feature_extractor = None
+        self.image_path = ""
+
+    def load_model(self) :
+
+        from transformers import ViTFeatureExtractor, ViTForImageClassification
+
+        self.model =  ViTForImageClassification.from_pretrained('google/vit-base-patch16-224', 
+                                                #   num_labels=100,
+                                                #   ignore_mismatched_sizes=True,
+                                                #   map_location=torch.device('cpu') 
+                                                  )
+
+        self.feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224')
+
+
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.model.to(self.device)
+        self.model.eval()
+
+    def set_image_path(self, image_path):
+        self.image_path = image_path
+
+    def run(self):
+        # Load and process the image
+        if not self.model:
+            self.load_model()
+
+        if self.image_path:
+            image = Image.open(self.image_path).convert('RGB').resize((244,244))
+            processed_image = self.feature_extractor(images=image, return_tensors="pt")
+            processed_image = {k: v.to(self.device) for k, v in processed_image.items()}
+            with torch.no_grad():
+                outputs = self.model(**processed_image)
+
+            predicted_class_idx = torch.argmax(outputs.logits, dim=-1).item()
+            predicted_class =  self.model.config.id2label[predicted_class_idx]
+
+            self.result_signal.emit(predicted_class)
+        self.finished_signal.emit()