Update app.py

app.py

@@ -9,9 +9,9 @@ import logging
 logging.basicConfig(level=logging.DEBUG)
 
 # Load the pre-trained model and feature extractor
-model_name = "JoshuaKelleyDs/quickdraw-DeiT-tiny-finetune"
+model_name = "JoshuaKelleyDs/quickdraw-MobileVITV2-2.0-Finetune"
 logging.info("Loading image processor and model...")
-feature_extractor = AutoFeatureExtractor.from_pretrained(model_name, image_mean=[0.5], image_std=[0.5])
+feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
 model = AutoModelForImageClassification.from_pretrained(model_name)
 
 # Define the prediction function
@@ -27,13 +27,8 @@ def predict(image):
         logging.debug("Converting to NumPy array...")
         image = np.array(image).astype('uint8')
         logging.debug("Converting NumPy array to PIL image...")
-        image = Image.fromarray(image, 'RGBA').convert('L')  # Convert to grayscale (L mode)
+        image = Image.fromarray(image, 'RGBA').convert('RGB')
         logging.debug("Image converted successfully.")
-        
-        # Ensure the image has the correct number of dimensions
-        if len(image.size) == 2:
-            image = np.expand_dims(image, axis=-1)
-            logging.debug("Added dimension to image to match model input requirements.")
 
         logging.info("Processing image...")
         inputs = feature_extractor(images=image, return_tensors="pt")