Update app.py

app.py

@@ -1,16 +1,15 @@
 import gradio as gr
-from transformers import pipeline, AutoFeatureExtractor, SegformerForSemanticSegmentation
+from transformers import pipeline, AutoFeatureExtractor, AutoModelForImageClassification
 
+extractor = AutoFeatureExtractor.from_pretrained("nateraw/vit-age-classifier")
 
-extractor = AutoFeatureExtractor.from_pretrained("nickmuchi/segformer-b4-finetuned-segments-sidewalk")
+model = AutoModelForImageClassification.from_pretrained("nateraw/vit-age-classifier")
 
-model = SegformerForSemanticSegmentation.from_pretrained("nickmuchi/segformer-b4-finetuned-segments-sidewalk")
-model.eval()
-pipe = pipeline(task="image-segmentation", 
+pipe = pipeline(task="image-classification", 
                 model=model,
-                feature_extractor=extractor)
+                feature_extractor=extractor,)
 
 gr.Interface.load(pipe, 
-                    title="Segformer",
+                    title="age class",
                     description="Add description",
                     ).launch()