Add Vision Transformer demo for image classification

Author: dhruvidave348

computer_vision/vision_transformer.py

@@ -4,24 +4,15 @@
 
 Classify images using a pretrained Vision Transformer (ViT)
 from Hugging Face Transformers.
-
-Can be used as a demo or imported in other scripts.
-
-Source:
-https://huggingface.co/docs/transformers/model_doc/vit
 """
 
-try:
-    import requests
-    import torch
-    from io import BytesIO
-    from PIL import Image
-    from transformers import ViTForImageClassification, ViTImageProcessor
-except ImportError as e:
-    raise ImportError(
-        "This module requires 'torch', 'transformers', 'PIL', and 'requests'. "
-        "Install them with: pip install torch transformers pillow requests"
-    ) from e
+from io import BytesIO
+from typing import Optional
+
+import requests
+import torch
+from PIL import Image, UnidentifiedImageError
+from transformers import ViTForImageClassification, ViTImageProcessor
 
 
 def classify_image(image: Image.Image) -> str:
@@ -39,21 +30,23 @@ def classify_image(image: Image.Image) -> str:
     return model.config.id2label[predicted_class_idx]
 
 
-def demo(url: str = None) -> None:
+def demo(url: Optional[str] = None) -> None:
     """
     Run a demo using a sample image or provided URL.
-    
+
     Args:
-        url (str): URL of the image. If None, uses a default cat image.
+        url (Optional[str]): URL of the image. If None, uses default cat image.
     """
     if url is None:
-        url = "https://images.unsplash.com/photo-1592194996308-7b43878e84a6"  # default example image
+        url = (
+            "https://images.unsplash.com/photo-1592194996308-7b43878e84a6"
+        )  # default example image
 
     try:
         response = requests.get(url, timeout=10)
         response.raise_for_status()
         image = Image.open(BytesIO(response.content))
-    except Exception as e:
+    except (requests.RequestException, UnidentifiedImageError) as e:
         print(f"Failed to load image from {url}. Error: {e}")
         return