Update app.py

app.py

@@ -1,32 +1,29 @@
 import gradio as gr
 import numpy as np
-import torch
 from PIL import Image, ImageDraw
-from transformers import SegformerForSemanticSegmentation
-from torchvision.transforms import Resize, ToTensor, Normalize
+from transformers import TFSegformerForSemanticSegmentation
+import tensorflow as tf
 
 # Load models from Hugging Face
-part_seg_model = SegformerForSemanticSegmentation.from_pretrained("Mohaddz/huggingCars")
-damage_seg_model = SegformerForSemanticSegmentation.from_pretrained("Mohaddz/DamageSegMohaddz/DamageSeg")
+part_seg_model = TFSegformerForSemanticSegmentation.from_pretrained("Mohaddz/huggingCars")
+damage_seg_model = TFSegformerForSemanticSegmentation.from_pretrained("Mohaddz/DamageSegMohaddz/DamageSeg")
 
 # Define your labels
 part_labels = ["front-bumper", "fender", "hood", "door", "trunk", "cars-8I1q"]  # Add all your part labels
 damage_labels = ["dent", "scratch", "misalignment", "crack", "etc"]  # Add all your damage labels
 
 def preprocess_image(image):
-    # Resize and normalize the 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
+    # Resize the image
+    image = tf.image.resize(image, (512, 512))
+    # Normalize the image
+    image = tf.keras.applications.imagenet_utils.preprocess_input(image)
+    return tf.expand_dims(image, 0)  # Add batch dimension
 
 def inference_seg(model, image):
-    with torch.no_grad():
-        outputs = model(image)
+    outputs = model(image, training=False)
     logits = outputs.logits
-    mask = torch.argmax(logits, dim=1).squeeze().numpy()
-    return mask
+    mask = tf.argmax(logits, axis=-1)
+    return mask.numpy().squeeze()
 
 def inference_part_seg(image):
     preprocessed_image = preprocess_image(image)
@@ -62,7 +59,7 @@ def create_one_hot_vector(part_damage_pairs):
     return vector
 
 def visualize_results(image, part_mask, damage_mask):
-    img = Image.fromarray(image)
+    img = Image.fromarray((image * 255).astype('uint8'))
     draw = ImageDraw.Draw(img)
     
     for i in range(img.width):