Update app.py

app.py

@@ -1,12 +1,13 @@
 import gradio as gr
 import numpy as np
-import tensorflow as tf
+import torch
 from PIL import Image, ImageDraw
-from transformers import TFSegformerForSemanticSegmentation
+from transformers import SegformerForSemanticSegmentation
+from torchvision.transforms import Resize, ToTensor, Normalize
 
 # Load models from Hugging Face
-part_seg_model = TFSegformerForSemanticSegmentation.from_pretrained("Mohaddz/huggingCars")
-damage_seg_model = TFSegformerForSemanticSegmentation.from_pretrained("Mohaddz/DamageSeg")
+part_seg_model = SegformerForSemanticSegmentation.from_pretrained("Mohaddz/huggingCars")
+damage_seg_model = SegformerForSemanticSegmentation.from_pretrained("Mohaddz/DamageSegMohaddz/DamageSeg")
 
 # Define your labels
 part_labels = ["front-bumper", "fender", "hood", "door", "trunk", "cars-8I1q"]  # Add all your part labels
@@ -14,23 +15,26 @@ damage_labels = ["dent", "scratch", "misalignment", "crack", "etc"]  # Add all y
 
 def preprocess_image(image):
     # Resize and normalize the image
-    image = tf.image.resize(image, (512, 512))
-    image = tf.keras.applications.imagenet_utils.preprocess_input(image)
-    return image
+    transform = Resize((512, 512))
+    image = transform(Image.fromarray(image))
+    image = ToTensor()(image)
+    image = Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])(image)
+    return image.unsqueeze(0)  # Add batch dimension
+
+def inference_seg(model, image):
+    with torch.no_grad():
+        outputs = model(image)
+    logits = outputs.logits
+    mask = torch.argmax(logits, dim=1).squeeze().numpy()
+    return mask
 
 def inference_part_seg(image):
     preprocessed_image = preprocess_image(image)
-    outputs = part_seg_model(preprocessed_image, training=False)
-    logits = outputs.logits
-    mask = tf.argmax(logits, axis=-1)
-    return mask.numpy()
+    return inference_seg(part_seg_model, preprocessed_image)
 
 def inference_damage_seg(image):
     preprocessed_image = preprocess_image(image)
-    outputs = damage_seg_model(preprocessed_image, training=False)
-    logits = outputs.logits
-    mask = tf.argmax(logits, axis=-1)
-    return mask.numpy()
+    return inference_seg(damage_seg_model, preprocessed_image)
 
 def combine_masks(part_mask, damage_mask):
     part_damage_pairs = []
@@ -58,7 +62,7 @@ def create_one_hot_vector(part_damage_pairs):
     return vector
 
 def visualize_results(image, part_mask, damage_mask):
-    img = Image.fromarray((image * 255).astype(np.uint8))
+    img = Image.fromarray(image)
     draw = ImageDraw.Draw(img)
     
     for i in range(img.width):
@@ -71,10 +75,6 @@ def visualize_results(image, part_mask, damage_mask):
     return img
 
 def process_image(image):
-    # Convert to numpy array if it's not already
-    if isinstance(image, Image.Image):
-        image = np.array(image)
-    
     # Perform inference
     part_mask = inference_part_seg(image)
     damage_mask = inference_damage_seg(image)
@@ -100,7 +100,7 @@ def gradio_interface(input_image):
 
 iface = gr.Interface(
     fn=gradio_interface,
-    inputs=gr.Image(type="pil"),
+    inputs=gr.Image(type="numpy"),
     outputs=[
         gr.Image(type="pil", label="Detected Damage"),
         gr.Textbox(label="Damage Description"),