Update app.py

app.py

@@ -2,25 +2,35 @@ import gradio as gr
 import numpy as np
 import tensorflow as tf
 from PIL import Image, ImageDraw
-import io
+from transformers import TFSegformerForSemanticSegmentation
 
-# Load your models here
- part_seg_model = tf.keras.models.load_model('path_to_part_seg_model')
- damage_seg_model = tf.keras.models.load_model('path_to_damage_seg_model')
+# Load models from Hugging Face
+part_seg_model = TFSegformerForSemanticSegmentation.from_pretrained("Mohaddz/huggingCars")
+damage_seg_model = TFSegformerForSemanticSegmentation.from_pretrained("Mohaddz/DamageSeg")
 
 # Define your labels
-part_labels = ["front-bumper", "fender", "hood", "door", "trunk"]  # Add all your part labels
-damage_labels = ["dent", "scratch", "misalignment", "crack"]  # Add all your damage 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
+    image = tf.image.resize(image, (512, 512))
+    image = tf.keras.applications.imagenet_utils.preprocess_input(image)
+    return image
 
 def inference_part_seg(image):
-    # Implement your part segmentation inference here
-    # For now, we'll return a random mask
-    return np.random.randint(0, len(part_labels), size=image.shape[:2])
+    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()
 
 def inference_damage_seg(image):
-    # Implement your damage segmentation inference here
-    # For now, we'll return a random mask
-    return np.random.randint(0, len(damage_labels), size=image.shape[:2])
+    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()
 
 def combine_masks(part_mask, damage_mask):
     part_damage_pairs = []
@@ -48,7 +58,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(np.uint8))
     draw = ImageDraw.Draw(img)
     
     for i in range(img.width):