Update app.py

app.py

@@ -1,3 +1,32 @@
 import gradio as gr
+from transformers import AutoFeatureExtractor, SegformerForSemanticSegmentation
+from PIL import Image
+import requests
+import matplotlib.pyplot as plt
+import torch.nn as nn
 
-gr.Interface.load("models/mattmdjaga/segformer_b2_clothes").launch()
+extractor = AutoFeatureExtractor.from_pretrained("mattmdjaga/segformer_b2_clothes")
+model = SegformerForSemanticSegmentation.from_pretrained("mattmdjaga/segformer_b2_clothes")
+
+def predict(inp):
+    inputs = extractor(images=inp, return_tensors="pt")
+    outputs = model(**inputs)
+    logits = outputs.logits.cpu()
+    upsampled_logits = nn.functional.interpolate(
+        logits,
+        size=image.size[::-1],
+        mode="bilinear",
+        align_corners=False,
+    )
+    pred_seg = upsampled_logits.argmax(dim=1)[0]
+    pred_seg[pred_seg != 4] = 0
+    arr_seg = pred_seg.cpu().numpy().astype("uint8")
+    arr_seg *= 255
+    pil_seg = Image.fromarray(arr_seg)
+  
+    return pil_seg
+
+gr.Interface(fn=predict, 
+             inputs=gr.Image(type="pil"),
+             outputs="image",
+             ).launch()