fetched and marged changes made in the remote repo

.gitattributes

@@ -35,3 +35,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 models/*.safetensors filter=lfs diff=lfs merge=lfs -text
 corrected[[:space:]]model[[:space:]]path filter=lfs diff=lfs merge=lfs -text
+*.ipynb filter=lfs diff=lfs merge=lfs -text

app.py

@@ -5,28 +5,28 @@ import matplotlib.pyplot as plt
 import io
 from model import load_model, get_val_transform  # Import functions from model.py
 import numpy as np
+<<<<<<< HEAD
+=======
+
+# Load the model on GPU if available
+model = load_model(device=0 if torch.cuda.is_available() else -1)
+>>>>>>> 16e22fbe27ebb36b6090c462c63a4d127310b2b8
 
-# Load model and image processor
-model, image_processor = load_model()
 val_transform = get_val_transform()
 
 # Define colors for bounding boxes
 COLORS = [[0.000, 0.447, 0.741], [0.850, 0.325, 0.098], [0.929, 0.694, 0.125],
           [0.494, 0.184, 0.556], [0.466, 0.674, 0.188], [0.301, 0.745, 0.933]]
 
-def preprocess(image):
-    """Preprocess image using the validation transform."""
-    numpy_image = np.array(image)
-    transformed = val_transform(image=numpy_image, category=[])  # No bounding boxes needed for input
-    return transformed["image"]
-
-def postprocess(image, outputs, threshold):
-    """Postprocess outputs to draw bounding boxes on the image."""
+def get_output_figure(pil_img, results, threshold):
     plt.figure(figsize=(12, 8))
-    plt.imshow(image)
+    plt.imshow(pil_img)
     ax = plt.gca()
 
-    for box, label, score in zip(outputs['boxes'], outputs['labels'], outputs['scores']):
+    for result in results:
+        score = result['score']
+        label = result['label']
+        box = list(result['box'].values())
         if score > threshold:
             color = COLORS[hash(label) % len(COLORS)]
             ax.add_patch(
@@ -40,10 +40,8 @@ def postprocess(image, outputs, threshold):
                 box[0], box[1] - 5, text, fontsize=10,
                 bbox=dict(facecolor='yellow', alpha=0.5, edgecolor='none')
             )
-
     plt.axis('off')
 
-    # Convert matplotlib figure to PIL image
     buf = io.BytesIO()
     plt.savefig(buf, bbox_inches='tight', dpi=100)
     buf.seek(0)
@@ -51,25 +49,9 @@ def postprocess(image, outputs, threshold):
     return Image.open(buf)
 
 def detect(image, threshold=0.5):
-    """Run the detection pipeline."""
-    # Preprocess the image
-    processed_image = preprocess(image)
-
-    # Convert to tensor for the model
-    inputs = image_processor(images=processed_image, return_tensors="pt")
-
-    # Run the model
-    outputs = model(**inputs)
-
-    # Convert the outputs to a more usable format
-    results = {
-        "boxes": outputs.logits.argmax(dim=-1).tolist(),  # Replace with actual box extraction logic
-        "labels": outputs.logits.argmax(dim=-1).tolist(),  # Replace with actual label extraction logic
-        "scores": outputs.scores.tolist(),  # Replace with actual score extraction logic
-    }
-
-    # Postprocess and return the annotated image
-    return postprocess(image, results, threshold)
+    results = model(image)
+    output_image = get_output_figure(image, results, threshold)
+    return output_image
 
 # Build the Gradio app
 with gr.Blocks() as demo:
@@ -84,7 +66,7 @@ with gr.Blocks() as demo:
     with gr.Row():
         image_input = gr.Image(label="Input Image", type="pil")
         threshold_slider = gr.Slider(
-            minimum=0.0, maximum=1.0, step=0.05, value=0.7, label="Detection Threshold"
+            minimum=0.0, maximum=1.0, step=0.05, value=0.5, label="Detection Threshold"
         )
 
     output_image = gr.Image(label="Output Prediction", type="pil")

detr_fine_tuning_custom_dataset.ipynb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b92462743cc8ab359bd1a790f7213b6494fb0726cbf70c25982060ff75e7b06e
+size 7982824

model.py

@@ -1,5 +1,6 @@
 import albumentations as A
 from transformers import AutoModelForObjectDetection, AutoImageProcessor
+from transformers import pipeline
 
 
 # Mapping for labels and IDs
@@ -52,11 +53,15 @@ def get_val_transform():
     )
 
 # Load the model
-def load_model():
-    model_path = "model"  # Path to your saved model
-    image_processor = AutoImageProcessor.from_pretrained(model_path)
-    model = AutoModelForObjectDetection.from_pretrained(
-        model_path,  # This will automatically use model.safetensors
-        ignore_mismatched_sizes=True
+def load_model(device: int = -1):
+    """
+    Load the DETR model pipeline.
+    :param device: Specify device to load the model (-1 for CPU, 0 for GPU).
+    :return: Hugging Face object-detection pipeline.
+    """
+    model_pipeline = pipeline(
+        "object-detection", 
+        model="sergiopaniego/detr-resnet-50-dc5-fashionpedia-finetuned", 
+        device=device
     )
-    return model, image_processor
+    return model_pipeline