Upload 15 files

app.py

@@ -0,0 +1,75 @@
+### 1. Imports and class names setup ###
+import gradio as gr
+import os
+import torch
+
+from class_names import class_names
+from model import Load_model
+from timeit import default_timer as timer
+from typing import Tuple, Dict
+
+
+### 1. Model and transforms preparation ###
+
+# Create model and transform
+model, transforms = Load_model()
+
+# Load saved weights
+def load_checkpoint(checkpoint_file, model, device='cpu'):
+    print("=> Loading checkpoint")
+    checkpoint = torch.load(checkpoint_file, map_location=device)
+    model.load_state_dict(checkpoint["state_dict"])
+load_checkpoint('model_checkpoint.pt', model)
+
+### 2. Predict function ###
+
+# Create predict function
+def predict(img) -> Tuple[Dict, float]:
+    """Transforms and performs a prediction on img and returns prediction and time taken.
+    """
+    # Start the timer
+    start_time = timer()
+
+    # Transform the target image and add a batch dimension
+    img = transforms(img).unsqueeze(0)
+
+    # Put model into evaluation mode and turn on inference mode
+    model.eval()
+    with torch.inference_mode():
+        # Pass the transformed image through the model and turn the prediction logits into prediction probabilities
+        pred_probs = torch.softmax(model(img), dim=1)
+
+    # Create a prediction label and prediction probability dictionary for each prediction class (this is the required format for Gradio's output parameter)
+    pred_labels_and_probs = {class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))}
+
+    # Calculate the prediction time
+    pred_time = round(timer() - start_time, 5)
+
+    # Return the prediction dictionary and prediction time
+    return pred_labels_and_probs, pred_time
+
+
+### 3. Gradio app ###
+
+# Create title, description and article strings
+title = "BirdVision 500 🦅🦆🐦🕊🦤🦢🦜"
+description = "A model based on YoLov8 classification 500 birds."
+article = "Created on [GITHUB](https://github.com/vvduc1803?tab=repositories/)."
+
+# Create examples list from "examples/" directory
+example_list = [["examples/" + example] for example in os.listdir("examples")]
+
+# Create the Gradio demo
+demo = gr.Interface(fn=predict,  # mapping function from input to output
+                    inputs=gr.Image(type="pil"),  # what are the inputs?
+                    outputs=[gr.Label(num_top_classes=10, label="Predictions"),  # what are the outputs?
+                             gr.Number(label="Prediction time (s)")],
+                    # our fn has two outputs, therefore we have two outputs
+                    # Create examples list from "examples/" directory
+                    examples=example_list,
+                    title=title,
+                    description=description,
+                    article=article)
+
+# Launch the demo!
+demo.launch()

class_names.py

@@ -0,0 +1,117 @@
+class_names = ['ABBOTTS BABBLER', 'ABBOTTS BOOBY', 'ABYSSINIAN GROUND HORNBILL', 'AFRICAN CROWNED CRANE', 'AFRICAN '
+                                                                                                          'EMERALD '
+                                                                                                          'CUCKOO',
+               'AFRICAN FIREFINCH', 'AFRICAN OYSTER CATCHER', 'AFRICAN PIED HORNBILL', 'AFRICAN PYGMY GOOSE',
+               'ALBATROSS', 'ALBERTS TOWHEE', 'ALEXANDRINE PARAKEET', 'ALPINE CHOUGH', 'ALTAMIRA YELLOWTHROAT',
+               'AMERICAN AVOCET', 'AMERICAN BITTERN', 'AMERICAN COOT', 'AMERICAN FLAMINGO', 'AMERICAN GOLDFINCH',
+               'AMERICAN KESTREL', 'AMERICAN PIPIT', 'AMERICAN REDSTART', 'AMERICAN ROBIN', 'AMERICAN WIGEON',
+               'AMETHYST WOODSTAR', 'ANDEAN GOOSE', 'ANDEAN LAPWING', 'ANDEAN SISKIN', 'ANHINGA', 'ANIANIAU',
+               'ANNAS HUMMINGBIRD', 'ANTBIRD', 'ANTILLEAN EUPHONIA', 'APAPANE', 'APOSTLEBIRD', 'ARARIPE MANAKIN',
+               'ASHY STORM PETREL', 'ASHY THRUSHBIRD', 'ASIAN CRESTED IBIS', 'ASIAN DOLLARD BIRD', 'AUCKLAND SHAQ',
+               'AUSTRAL CANASTERO', 'AUSTRALASIAN FIGBIRD', 'AVADAVAT', 'AZARAS SPINETAIL', 'AZURE BREASTED PITTA',
+               'AZURE JAY', 'AZURE TANAGER', 'AZURE TIT', 'BAIKAL TEAL', 'BALD EAGLE', 'BALD IBIS', 'BALI STARLING',
+               'BALTIMORE ORIOLE', 'BANANAQUIT', 'BAND TAILED GUAN', 'BANDED BROADBILL', 'BANDED PITA',
+               'BANDED STILT', 'BAR-TAILED GODWIT', 'BARN OWL', 'BARN SWALLOW', 'BARRED PUFFBIRD',
+               'BARROWS GOLDENEYE', 'BAY-BREASTED WARBLER', 'BEARDED BARBET', 'BEARDED BELLBIRD', 'BEARDED REEDLING',
+               'BELTED KINGFISHER', 'BIRD OF PARADISE', 'BLACK AND YELLOW BROADBILL', 'BLACK BAZA', 'BLACK COCKATO',
+               'BLACK FACED SPOONBILL', 'BLACK FRANCOLIN', 'BLACK HEADED CAIQUE', 'BLACK NECKED STILT',
+               'BLACK SKIMMER', 'BLACK SWAN', 'BLACK TAIL CRAKE', 'BLACK THROATED BUSHTIT', 'BLACK THROATED HUET',
+               'BLACK THROATED WARBLER', 'BLACK VENTED SHEARWATER', 'BLACK VULTURE', 'BLACK-CAPPED CHICKADEE',
+               'BLACK-NECKED GREBE', 'BLACK-THROATED SPARROW', 'BLACKBURNIAM WARBLER', 'BLONDE CRESTED WOODPECKER',
+               'BLOOD PHEASANT', 'BLUE COAU', 'BLUE DACNIS', 'BLUE GRAY GNATCATCHER', 'BLUE GROSBEAK', 'BLUE GROUSE',
+               'BLUE HERON', 'BLUE MALKOHA', 'BLUE THROATED TOUCANET', 'BOBOLINK', 'BORNEAN BRISTLEHEAD',
+               'BORNEAN LEAFBIRD', 'BORNEAN PHEASANT', 'BRANDT CORMARANT', 'BREWERS BLACKBIRD', 'BROWN CREPPER',
+               'BROWN HEADED COWBIRD', 'BROWN NOODY', 'BROWN THRASHER', 'BUFFLEHEAD', 'BULWERS PHEASANT', 'BURCHELLS '
+                                                                                                          'COURSER',
+               'BUSH TURKEY', 'CAATINGA CACHOLOTE', 'CACTUS WREN', 'CALIFORNIA CONDOR', 'CALIFORNIA GULL',
+               'CALIFORNIA QUAIL', 'CAMPO FLICKER', 'CANARY', 'CANVASBACK', 'CAPE GLOSSY STARLING', 'CAPE LONGCLAW',
+               'CAPE MAY WARBLER', 'CAPE ROCK THRUSH', 'CAPPED HERON', 'CAPUCHINBIRD', 'CARMINE BEE-EATER',
+               'CASPIAN TERN', 'CASSOWARY', 'CEDAR WAXWING', 'CERULEAN WARBLER', 'CHARA DE COLLAR', 'CHATTERING '
+                                                                                                    'LORY',
+               'CHESTNET BELLIED EUPHONIA', 'CHINESE BAMBOO PARTRIDGE', 'CHINESE POND HERON', 'CHIPPING SPARROW',
+               'CHUCAO TAPACULO', 'CHUKAR PARTRIDGE', 'CINNAMON ATTILA', 'CINNAMON FLYCATCHER', 'CINNAMON TEAL',
+               'CLARKS GREBE', 'CLARKS NUTCRACKER', 'COCK OF THE  ROCK', 'COCKATOO', 'COLLARED ARACARI',
+               'COLLARED CRESCENTCHEST', 'COMMON FIRECREST', 'COMMON GRACKLE', 'COMMON HOUSE MARTIN', 'COMMON IORA',
+               'COMMON LOON', 'COMMON POORWILL', 'COMMON STARLING', 'COPPERY TAILED COUCAL', 'CRAB PLOVER',
+               'CRANE HAWK', 'CREAM COLORED WOODPECKER', 'CRESTED AUKLET', 'CRESTED CARACARA', 'CRESTED COUA',
+               'CRESTED FIREBACK', 'CRESTED KINGFISHER', 'CRESTED NUTHATCH', 'CRESTED OROPENDOLA', 'CRESTED SERPENT '
+                                                                                                   'EAGLE',
+               'CRESTED SHRIKETIT', 'CRESTED WOOD PARTRIDGE', 'CRIMSON CHAT', 'CRIMSON SUNBIRD', 'CROW',
+               'CROWNED PIGEON', 'CUBAN TODY', 'CUBAN TROGON', 'CURL CRESTED ARACURI', 'D-ARNAUDS BARBET',
+               'DALMATIAN PELICAN', 'DARJEELING WOODPECKER', 'DARK EYED JUNCO', 'DAURIAN REDSTART', 'DEMOISELLE '
+                                                                                                    'CRANE',
+               'DOUBLE BARRED FINCH', 'DOUBLE BRESTED CORMARANT', 'DOUBLE EYED FIG PARROT', 'DOWNY WOODPECKER',
+               'DUSKY LORY', 'DUSKY ROBIN', 'EARED PITA', 'EASTERN BLUEBIRD', 'EASTERN BLUEBONNET', 'EASTERN GOLDEN '
+                                                                                                    'WEAVER',
+               'EASTERN MEADOWLARK', 'EASTERN ROSELLA', 'EASTERN TOWEE', 'EASTERN WIP POOR WILL', 'EASTERN YELLOW '
+                                                                                                  'ROBIN',
+               'ECUADORIAN HILLSTAR', 'EGYPTIAN GOOSE', 'ELEGANT TROGON', 'ELLIOTS  PHEASANT', 'EMERALD TANAGER',
+               'EMPEROR PENGUIN', 'EMU', 'ENGGANO MYNA', 'EURASIAN BULLFINCH', 'EURASIAN GOLDEN ORIOLE',
+               'EURASIAN MAGPIE', 'EUROPEAN GOLDFINCH', 'EUROPEAN TURTLE DOVE', 'EVENING GROSBEAK', 'FAIRY BLUEBIRD',
+               'FAIRY PENGUIN', 'FAIRY TERN', 'FAN TAILED WIDOW', 'FASCIATED WREN', 'FIERY MINIVET', 'FIORDLAND '
+                                                                                                     'PENGUIN',
+               'FIRE TAILLED MYZORNIS', 'FLAME BOWERBIRD', 'FLAME TANAGER', 'FRIGATE', 'FRILL BACK PIGEON',
+               'GAMBELS QUAIL', 'GANG GANG COCKATOO', 'GILA WOODPECKER', 'GILDED FLICKER', 'GLOSSY IBIS',
+               'GO AWAY BIRD', 'GOLD WING WARBLER', 'GOLDEN BOWER BIRD', 'GOLDEN CHEEKED WARBLER',
+               'GOLDEN CHLOROPHONIA', 'GOLDEN EAGLE', 'GOLDEN PARAKEET', 'GOLDEN PHEASANT', 'GOLDEN PIPIT',
+               'GOULDIAN FINCH', 'GRANDALA', 'GRAY CATBIRD', 'GRAY KINGBIRD', 'GRAY PARTRIDGE', 'GREAT ARGUS',
+               'GREAT GRAY OWL', 'GREAT JACAMAR', 'GREAT KISKADEE', 'GREAT POTOO', 'GREAT TINAMOU', 'GREAT XENOPS',
+               'GREATER PEWEE', 'GREATER PRAIRIE CHICKEN', 'GREATOR SAGE GROUSE', 'GREEN BROADBILL', 'GREEN JAY',
+               'GREEN MAGPIE', 'GREEN WINGED DOVE', 'GREY CUCKOOSHRIKE', 'GREY HEADED FISH EAGLE', 'GREY PLOVER',
+               'GROVED BILLED ANI', 'GUINEA TURACO', 'GUINEAFOWL', 'GURNEYS PITTA', 'GYRFALCON', 'HAMERKOP',
+               'HARLEQUIN DUCK', 'HARLEQUIN QUAIL', 'HARPY EAGLE', 'HAWAIIAN GOOSE', 'HAWFINCH', 'HELMET VANGA',
+               'HEPATIC TANAGER', 'HIMALAYAN BLUETAIL', 'HIMALAYAN MONAL', 'HOATZIN', 'HOODED MERGANSER', 'HOOPOES',
+               'HORNED GUAN', 'HORNED LARK', 'HORNED SUNGEM', 'HOUSE FINCH', 'HOUSE SPARROW', 'HYACINTH MACAW',
+               'IBERIAN MAGPIE', 'IBISBILL', 'IMPERIAL SHAQ', 'INCA TERN', 'INDIAN BUSTARD', 'INDIAN PITTA',
+               'INDIAN ROLLER', 'INDIAN VULTURE', 'INDIGO BUNTING', 'INDIGO FLYCATCHER', 'INLAND DOTTEREL',
+               'IVORY BILLED ARACARI', 'IVORY GULL', 'IWI', 'JABIRU', 'JACK SNIPE', 'JACOBIN PIGEON',
+               'JANDAYA PARAKEET', 'JAPANESE ROBIN', 'JAVA SPARROW', 'JOCOTOCO ANTPITTA', 'KAGU', 'KAKAPO',
+               'KILLDEAR', 'KING EIDER', 'KING VULTURE', 'KIWI', 'KOOKABURRA', 'LARK BUNTING', 'LAUGHING GULL',
+               'LAZULI BUNTING', 'LESSER ADJUTANT', 'LILAC ROLLER', 'LIMPKIN', 'LITTLE AUK', 'LOGGERHEAD SHRIKE',
+               'LONG-EARED OWL', 'LOONEY BIRDS', 'LUCIFER HUMMINGBIRD', 'MAGPIE GOOSE', 'MALABAR HORNBILL',
+               'MALACHITE KINGFISHER', 'MALAGASY WHITE EYE', 'MALEO', 'MALLARD DUCK', 'MANDRIN DUCK',
+               'MANGROVE CUCKOO', 'MARABOU STORK', 'MASKED BOBWHITE', 'MASKED BOOBY', 'MASKED LAPWING',
+               'MCKAYS BUNTING', 'MERLIN', 'MIKADO  PHEASANT', 'MILITARY MACAW', 'MOURNING DOVE', 'MYNA',
+               'NICOBAR PIGEON', 'NOISY FRIARBIRD', 'NORTHERN BEARDLESS TYRANNULET', 'NORTHERN CARDINAL',
+               'NORTHERN FLICKER', 'NORTHERN FULMAR', 'NORTHERN GANNET', 'NORTHERN GOSHAWK', 'NORTHERN JACANA',
+               'NORTHERN MOCKINGBIRD', 'NORTHERN PARULA', 'NORTHERN RED BISHOP', 'NORTHERN SHOVELER', 'OCELLATED '
+                                                                                                      'TURKEY',
+               'OKINAWA RAIL', 'ORANGE BRESTED BUNTING', 'ORIENTAL BAY OWL', 'ORNATE HAWK EAGLE', 'OSPREY',
+               'OSTRICH', 'OVENBIRD', 'OYSTER CATCHER', 'PAINTED BUNTING', 'PALILA', 'PALM NUT VULTURE',
+               'PARADISE TANAGER', 'PARAKETT  AKULET', 'PARUS MAJOR', 'PATAGONIAN SIERRA FINCH', 'PEACOCK',
+               'PEREGRINE FALCON', 'PHAINOPEPLA', 'PHILIPPINE EAGLE', 'PINK ROBIN', 'PLUSH CRESTED JAY',
+               'POMARINE JAEGER', 'PUFFIN', 'PUNA TEAL', 'PURPLE FINCH', 'PURPLE GALLINULE', 'PURPLE MARTIN',
+               'PURPLE SWAMPHEN', 'PYGMY KINGFISHER', 'PYRRHULOXIA', 'QUETZAL', 'RAINBOW LORIKEET', 'RAZORBILL',
+               'RED BEARDED BEE EATER', 'RED BELLIED PITTA', 'RED BILLED TROPICBIRD', 'RED BROWED FINCH', 'RED FACED '
+                                                                                                          'CORMORANT',
+               'RED FACED WARBLER', 'RED FODY', 'RED HEADED DUCK', 'RED HEADED WOODPECKER', 'RED KNOT', 'RED LEGGED '
+                                                                                                        'HONEYCREEPER'
+                                                                                                        '',
+               'RED NAPED TROGON', 'RED SHOULDERED HAWK', 'RED TAILED HAWK', 'RED TAILED THRUSH', 'RED WINGED '
+                                                                                                  'BLACKBIRD',
+               'RED WISKERED BULBUL', 'REGENT BOWERBIRD', 'RING-NECKED PHEASANT', 'ROADRUNNER', 'ROCK DOVE',
+               'ROSE BREASTED COCKATOO', 'ROSE BREASTED GROSBEAK', 'ROSEATE SPOONBILL', 'ROSY FACED LOVEBIRD',
+               'ROUGH LEG BUZZARD', 'ROYAL FLYCATCHER', 'RUBY CROWNED KINGLET', 'RUBY THROATED HUMMINGBIRD',
+               'RUDY KINGFISHER', 'RUFOUS KINGFISHER', 'RUFUOS MOTMOT', 'SAMATRAN THRUSH', 'SAND MARTIN',
+               'SANDHILL CRANE', 'SATYR TRAGOPAN', 'SAYS PHOEBE', 'SCARLET CROWNED FRUIT DOVE', 'SCARLET FACED '
+                                                                                                'LIOCICHLA',
+               'SCARLET IBIS', 'SCARLET MACAW', 'SCARLET TANAGER', 'SHOEBILL', 'SHORT BILLED DOWITCHER',
+               'SMITHS LONGSPUR', 'SNOW GOOSE', 'SNOWY EGRET', 'SNOWY OWL', 'SNOWY PLOVER', 'SORA', 'SPANGLED '
+                                                                                                    'COTINGA',
+               'SPLENDID WREN', 'SPOON BILED SANDPIPER', 'SPOTTED CATBIRD', 'SPOTTED WHISTLING DUCK', 'SRI LANKA BLUE '
+                                                                                                      'MAGPIE',
+               'STEAMER DUCK', 'STORK BILLED KINGFISHER', 'STRIATED CARACARA', 'STRIPED OWL', 'STRIPPED MANAKIN',
+               'STRIPPED SWALLOW', 'SUNBITTERN', 'SUPERB STARLING', 'SURF SCOTER', 'SWINHOES PHEASANT', 'TAILORBIRD',
+               'TAIWAN MAGPIE', 'TAKAHE', 'TASMANIAN HEN', 'TAWNY FROGMOUTH', 'TEAL DUCK', 'TIT MOUSE', 'TOUCHAN',
+               'TOWNSENDS WARBLER', 'TREE SWALLOW', 'TRICOLORED BLACKBIRD', 'TROPICAL KINGBIRD', 'TRUMPTER SWAN',
+               'TURKEY VULTURE', 'TURQUOISE MOTMOT', 'UMBRELLA BIRD', 'VARIED THRUSH', 'VEERY', 'VENEZUELIAN '
+                                                                                                'TROUPIAL', 'VERDIN',
+               'VERMILION FLYCATHER', 'VICTORIA CROWNED PIGEON', 'VIOLET BACKED STARLING', 'VIOLET GREEN SWALLOW',
+               'VIOLET TURACO', 'VULTURINE GUINEAFOWL', 'WALL CREAPER', 'WATTLED CURASSOW', 'WATTLED LAPWING',
+               'WHIMBREL', 'WHITE BROWED CRAKE', 'WHITE CHEEKED TURACO', 'WHITE CRESTED HORNBILL', 'WHITE EARED '
+                                                                                                   'HUMMINGBIRD',
+               'WHITE NECKED RAVEN', 'WHITE TAILED TROPIC', 'WHITE THROATED BEE EATER', 'WILD TURKEY',
+               'WILLOW PTARMIGAN', 'WILSONS BIRD OF PARADISE', 'WOOD DUCK', 'WOOD THRUSH', 'WRENTIT', 'YELLOW BELLIED '
+                                                                                                      'FLOWERPECKER',
+               'YELLOW CACIQUE', 'YELLOW HEADED BLACKBIRD']
+

model.py

@@ -0,0 +1,144 @@
+import torch
+import torch.nn as nn
+from torchvision import transforms
+
+class CNNBlock(nn.Module):
+    """Base block in CNN"""
+    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, bn_act=True):
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=not bn_act)
+        self.bn = nn.BatchNorm2d(out_channels)
+        self.silu = nn.SiLU()
+        self.use_bn_act = bn_act
+
+    def forward(self, x):
+        if self.use_bn_act:
+            x = self.silu(self.bn(self.conv(x)))
+
+            return x
+        else:
+            return self.conv(x)
+
+class BottleNeckBlock(nn.Module):
+    def __init__(self, channels, short_cut=True):
+        super().__init__()
+        self.short_cut = short_cut
+        self.Conv = nn.Sequential(CNNBlock(channels, channels//2, 3, 1, 1),
+                                  CNNBlock(channels//2, channels, 3, 1, 1))
+
+    def forward(self, x):
+        if self.short_cut:
+            return self.Conv(x) + x
+        else:
+            return self.Conv(x)
+
+class C2FBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, **kwargs):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.Conv = CNNBlock(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+        self.Conv_end = CNNBlock(int(0.5*(1+2)*out_channels), out_channels, kernel_size=1, stride=1, padding=0)
+        self.BottleNeck = BottleNeckBlock(out_channels//2, **kwargs)
+
+    def forward(self, x):
+        x = self.Conv(x)
+        x, x1 = torch.split(x, self.out_channels//2, dim=1)
+        x2 = self.BottleNeck(x1)
+        x = torch.cat([x, x1, x2], dim=1)
+        x = self.Conv_end(x)
+        return x
+
+class C2F_2_Block(nn.Module):
+    def __init__(self, in_channels, out_channels, **kwargs):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.Conv = CNNBlock(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+        self.Conv_end = CNNBlock(int(0.5*(2+2)*out_channels), out_channels, kernel_size=1, stride=1, padding=0)
+        self.BottleNeck = BottleNeckBlock(out_channels//2, **kwargs)
+
+    def forward(self, x):
+        x = self.Conv(x)
+        x, x1 = torch.split(x, self.out_channels//2, dim=1)
+        x2 = self.BottleNeck(x1)
+        x3 = self.BottleNeck(x2)
+        x = torch.cat([x, x1, x2, x3], dim=1)
+        x = self.Conv_end(x)
+        return x
+
+class SPPFBlock(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.Conv = CNNBlock(channels, channels, kernel_size=1, stride=1, padding=0)
+        self.Conv_end = CNNBlock(4*channels, channels, kernel_size=1, stride=1, padding=0)
+        self.MaxPool = nn.MaxPool2d(kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x):
+        x = self.Conv(x)
+        x = torch.cat([x, self.MaxPool(x), self.MaxPool(self.MaxPool(x)), self.MaxPool(self.MaxPool(self.MaxPool(x)))],
+                      dim=1)
+        x = self.Conv_end(x)
+        return x
+
+class Classifier(nn.Module):
+    def __init__(self, num_classes=500):
+        super().__init__()
+        self.Conv = nn.Sequential(CNNBlock(512, 1280, kernel_size=1, stride=1, padding=0))
+        self.Flatten = nn.Flatten()
+        self.Linear = nn.Sequential(nn.Linear(62720, num_classes))
+
+    def forward(self, x):
+        x = self.Conv(x)
+        x = self.Flatten(x)
+        x = self.Linear(x)
+        return x
+
+class Yolov8_cls(nn.Module):
+    """Model architecture based page: https://blog.roboflow.com/whats-new-in-yolov8/
+       and the ONNX file of yolov8_cls.onnx"""
+
+    def __init__(self, in_channels, num_classes=500):
+        super().__init__()
+        self.Block1 = nn.Sequential(CNNBlock(in_channels, 32, 3, 2, 1),
+                                    CNNBlock(32, 64, 3, 2, 1))
+
+        self.Block2 = C2FBlock(64, 64)
+
+        self.Block3 = nn.Sequential(CNNBlock(64, 128, 3, 2, 1),
+                                    C2F_2_Block(128, 128))
+
+        self.Block4 = nn.Sequential(CNNBlock(128, 256, 3, 2, 1),
+                                    C2F_2_Block(256, 256))
+
+        self.Block5 = nn.Sequential(CNNBlock(256, 512, 3, 2, 1),
+                                    C2F_2_Block(512, 512))
+
+        self.Block6 = Classifier(num_classes)
+
+    def forward(self, x):
+        x = self.Block1(x)
+        x = self.Block2(x)
+        x = self.Block3(x)
+        x = self.Block4(x)
+        x = self.Block5(x)
+        x = self.Block6(x)
+        return x
+
+
+def Load_model():
+    """Load model and transforms.
+    Returns:
+        model (torch.nn.Module): EffNetB2 feature extractor model.
+        transforms (torchvision.transforms): EffNetB2 image transforms.
+    """
+    IMAGE_SIZE= 224
+
+    model = Yolov8_cls(3)
+
+    transform = transforms.Compose([transforms.Resize(IMAGE_SIZE),
+                                    transforms.CenterCrop(IMAGE_SIZE),
+                                    transforms.ToTensor(),
+                                    transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])])
+
+    return model, transform

model_checkpoint.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3cc784f64ed37e9041760b8877c1e2bbb9ec9e46fdf05db1ed4441e66e6331ed
+size 425175633