app.py

import gradio as gr
import cv2
import torch
import torchvision.models as models
import torchvision.transforms as transforms
from PIL import Image
import numpy as np
import warnings

warnings.filterwarnings("ignore")


class ImageClassifier:
    def __init__(self, model_name):
        self.model = self.load_model(model_name)
        self.classes = self.load_classes()
        self.download_test_img()

    def download_test_img(self):
        # Images
        torch.hub.download_url_to_file(
            'https://user-images.githubusercontent.com/59380685/266264420-21575a83-4057-41cf-8a4a-b3ea6f332d79.jpg',
            'bus.jpg')
        torch.hub.download_url_to_file(
            'https://user-images.githubusercontent.com/59380685/266264536-82afdf58-6b9a-4568-b9df-551ee72cb6d9.jpg',
            'dogs.jpg')
        torch.hub.download_url_to_file(
            'https://user-images.githubusercontent.com/59380685/266264600-9d0c26ca-8ba6-45f2-b53b-4dc98460c43e.jpg',
            'zidane.jpg')

    def load_model(self, model_name):
        if model_name == "ResNet18":
            model = models.resnet18(pretrained=True)
        elif model_name == "AlexNet":
            model = models.alexnet(pretrained=True)
        elif model_name == "VGG16":
            model = models.vgg16(pretrained=True)
        elif model_name == "GoogLeNet":
            model = models.googlenet(pretrained=True)
        elif model_name == "ResNet50":
            model = models.resnet50(pretrained=True)
        elif model_name == "DenseNet121":
            model = models.densenet121(pretrained=True)
        elif model_name == "MobileNetV2":
            model = models.mobilenet_v2(pretrained=True)
        elif model_name == "ShuffleNetV2":
            model = models.shufflenet_v2_x1_0(pretrained=True)
        elif model_name == "SqueezeNet":
            model = models.squeezenet1_1(pretrained=True)
        elif model_name == "InceptionV3":
            model = models.inception_v3(pretrained=True)
        elif model_name == "ResNet101":
            model = models.resnet101(pretrained=True)
        elif model_name == "ResNet152":
            model = models.resnet152(pretrained=True)
        elif model_name == "WideResNet50":
            model = models.wide_resnet50_2(pretrained=True)
        elif model_name == "WideResNet101":
            model = models.wide_resnet101_2(pretrained=True)
        else:
            raise ValueError("Invalid model name")
        return model

    def load_classes(self):
        classes = ['0, tench', '1, goldfish', '2, great_white_shark', '3, tiger_shark', '4, hammerhead',
                   '5, electric_ray', '6, stingray',
                   '7, cock', '8, hen', '9, ostrich', '10, brambling', '11, goldfinch', '12, house_finch', '13, junco',
                   '14, indigo_bunting',
                   '15, robin', '16, bulbul', '17, jay', '18, magpie', '19, chickadee', '20, water_ouzel', '21, kite',
                   '22, bald_eagle',
                   '23, vulture', '24, great_grey_owl', '25, European_fire_salamander', '26, common_newt', '27, eft',
                   '28, spotted_salamander',
                   '29, axolotl', '30, bullfrog', '31, tree_frog', '32, tailed_frog', '33, loggerhead',
                   '34, leatherback_turtle', '35, mud_turtle',
                   '36, terrapin', '37, box_turtle', '38, banded_gecko', '39, common_iguana', '40, American_chameleon',
                   '41, whiptail',
                   '42, agama', '43, frilled_lizard', '44, alligator_lizard', '45, Gila_monster', '46, green_lizard',
                   '47, African_chameleon',
                   '48, Komodo_dragon', '49, African_crocodile', '50, American_alligator', '51, triceratops',
                   '52, thunder_snake', '53, ringneck_snake',
                   '54, hognose_snake', '55, green_snake', '56, king_snake', '57, garter_snake', '58, water_snake',
                   '59, vine_snake', '60, night_snake', '61, boa_constrictor', '62, rock_python', '63, Indian_cobra',
                   '64, green_mamba', '65, sea_snake',
                   '66, horned_viper', '67, diamondback', '68, sidewinder', '69, trilobite', '70, harvestman',
                   '71, scorpion',
                   '72, black_and_gold_garden_spider', '73, barn_spider', '74, garden_spider', '75, black_widow',
                   '76, tarantula',
                   '77, wolf_spider', '78, tick', '79, centipede', '80, black_grouse',
                   '81, ptarmigan', '82, ruffed_grouse', '83, prairie_chicken', '84, peacock', '85, quail',
                   '86, partridge', '87, African_grey',
                   '88, macaw', '89, sulphur-crested_cockatoo', '90, lorikeet', '91, coucal', '92, bee_eater',
                   '93, hornbill', '94, hummingbird', '95, jacamar', '96, toucan', '97, drake',
                   '98, red-breasted_merganser', '99, goose', '100, black_swan', '101, tusker', '102, echidna',
                   '103, platypus', '104, wallaby', '105, koala', '106, wombat', '107, jellyfish', '108, sea_anemone',
                   '109, brain_coral', '110, flatworm', '111, nematode', '112, conch', '113, snail', '114, slug',
                   '115, sea_slug', '116, chiton', '117, chambered_nautilus', '118, Dungeness_crab', '119, rock_crab',
                   '120, fiddler_crab', '121, king_crab', '122, American_lobster', '123, spiny_lobster',
                   '124, crayfish', '125, hermit_crab', '126, isopod', '127, white_stork', '128, black_stork',
                   '129, spoonbill', '130, flamingo', '131, little_blue_heron', '132, American_egret', '133, bittern',
                   '134, crane', '135, limpkin', '136, European_gallinule', '137, American_coot', '138, bustard',
                   '139, ruddy_turnstone', '140, red-backed_sandpiper', '141, redshank', '142, dowitcher',
                   '143, oystercatcher', '144, pelican', '145, king_penguin', '146, albatross', '147, grey_whale',
                   '148, killer_whale', '149, dugong', '150, sea_lion', '151, Chihuahua', '152, Japanese_spaniel',
                   '153, Maltese_dog', '154, Pekinese', '155, Shih-Tzu', '156, Blenheim_spaniel', '157, papillon',
                   '158, toy_terrier', '159, Rhodesian_ridgeback', '160, Afghan_hound', '161, basset', '162, beagle',
                   '163, bloodhound', '164, bluetick', '165, black-and-tan_coonhound', '166, Walker_hound',
                   '167, English_foxhound', '168, redbone', '169, borzoi', '170, Irish_wolfhound',
                   '171, Italian_greyhound', '172, whippet', '173, Ibizan_hound', '174, Norwegian_elkhound',
                   '175, otterhound', '176, Saluki', '177, Scottish_deerhound', '178, Weimaraner',
                   '179, Staffordshire_bullterrier', '180, American_Staffordshire_terrier', '181, Bedlington_terrier',
                   '182, Border_terrier', '183, Kerry_blue_terrier', '184, Irish_terrier', '185, Norfolk_terrier',
                   '186, Norwich_terrier', '187, Yorkshire_terrier', '188, wire-haired_fox_terrier',
                   '189, Lakeland_terrier', '190, Sealyham_terrier', '191, Airedale', '192, cairn',
                   '193, Australian_terrier', '194, Dandie_Dinmont', '195, Boston_bull', '196, miniature_schnauzer',
                   '197, giant_schnauzer', '198, standard_schnauzer', '199, Scotch_terrier', '200, Tibetan_terrier',
                   '201, silky_terrier', '202, soft-coated_wheaten_terrier', '203, West_Highland_white_terrier',
                   '204, Lhasa', '205, flat-coated_retriever', '206, curly-coated_retriever', '207, golden_retriever',
                   '208, Labrador_retriever', '209, Chesapeake_Bay_retriever', '210, German_short-haired_pointer',
                   '211, vizsla', '212, English_setter', '213, Irish_setter', '214, Gordon_setter',
                   '215, Brittany_spaniel', '216, clumber', '217, English_springer', '218, Welsh_springer_spaniel',
                   '219, cocker_spaniel', '220, Sussex_spaniel', '221, Irish_water_spaniel', '222, kuvasz',
                   '223, schipperke', '224, groenendael', '225, malinois', '226, briard', '227, kelpie',
                   '228, komondor', '229, Old_English_sheepdog', '230, Shetland_sheepdog', '231, collie',
                   '232, Border_collie', '233, Bouvier_des_Flandres', '234, Rottweiler', '235, German_shepherd',
                   '236, Doberman', '237, miniature_pinscher', '238, Greater_Swiss_Mountain_dog',
                   '239, Bernese_mountain_dog', '240, Appenzeller', '241, EntleBucher', '242, boxer',
                   '243, bull_mastiff', '244, Tibetan_mastiff', '245, French_bulldog', '246, Great_Dane',
                   '247, Saint_Bernard', '248, Eskimo_dog', '249, malamute', '250, Siberian_husky', '251, dalmatian',
                   '252, affenpinscher', '253, basenji', '254, pug', '255, Leonberg', '256, Newfoundland',
                   '257, Great_Pyrenees', '258, Samoyed', '259, Pomeranian', '260, chow', '261, keeshond',
                   '262, Brabancon_griffon', '263, Pembroke', '264, Cardigan', '265, toy_poodle',
                   '266, miniature_poodle', '267, standard_poodle', '268, Mexican_hairless', '269, timber_wolf',
                   '270, white_wolf', '271, red_wolf', '272, coyote', '273, dingo', '274, dhole',
                   '275, African_hunting_dog', '276, hyena', '277, red_fox', '278, kit_fox', '279, Arctic_fox',
                   '280, grey_fox', '281, tabby', '282, tiger_cat', '283, Persian_cat', '284, Siamese_cat',
                   '285, Egyptian_cat', '286, cougar', '287, lynx', '288, leopard', '289, snow_leopard', '290, jaguar',
                   '291, lion', '292, tiger', '293, cheetah', '294, brown_bear', '295, American_black_bear',
                   '296, ice_bear', '297, sloth_bear', '298, mongoose', '299, meerkat', '300, tiger_beetle',
                   '301, ladybug', '302, ground_beetle', '303, long-horned_beetle', '304, leaf_beetle',
                   '305, dung_beetle', '306, rhinoceros_beetle', '307, weevil', '308, fly', '309, bee', '310, ant',
                   '311, grasshopper', '312, cricket', '313, walking_stick', '314, cockroach', '315, mantis',
                   '316, cicada', '317, leafhopper', '318, lacewing', '319, dragonfly', '320, damselfly',
                   '321, admiral', '322, ringlet', '323, monarch', '324, cabbage_butterfly', '325, sulphur_butterfly',
                   '326, lycaenid', '327, starfish', '328, sea_urchin', '329, sea_cucumber', '330, wood_rabbit',
                   '331, hare', '332, Angora', '333, hamster', '334, porcupine', '335, fox_squirrel', '336, marmot',
                   '337, beaver', '338, guinea_pig', '339, sorrel', '340, zebra', '341, hog', '342, wild_boar',
                   '343, warthog', '344, hippopotamus', '345, ox', '346, water_buffalo', '347, bison', '348, ram',
                   '349, bighorn', '350, ibex', '351, hartebeest', '352, impala', '353, gazelle', '354, Arabian_camel',
                   '355, llama', '356, weasel', '357, mink', '358, polecat', '359, black-footed_ferret', '360, otter',
                   '361, skunk', '362, badger', '363, armadillo', '364, three-toed_sloth', '365, orangutan',
                   '366, gorilla', '367, chimpanzee', '368, gibbon', '369, siamang', '370, guenon', '371, patas',
                   '372, baboon', '373, macaque', '374, langur', '375, colobus', '376, proboscis_monkey',
                   '377, marmoset', '378, capuchin', '379, howler_monkey', '380, titi', '381, spider_monkey',
                   '382, squirrel_monkey', '383, Madagascar_cat', '384, indri', '385, Indian_elephant',
                   '386, African_elephant', '387, lesser_panda', '388, giant_panda', '389, barracouta', '390, eel',
                   '391, coho', '392, rock_beauty', '393, anemone_fish', '394, sturgeon', '395, gar', '396, lionfish',
                   '397, puffer', '398, abacus', '399, abaya', '400, academic_gown', '401, accordion',
                   '402, acoustic_guitar', '403, aircraft_carrier', '404, airliner', '405, airship', '406, altar',
                   '407, ambulance', '408, amphibian', '409, analog_clock', '410, apiary', '411, apron', '412, ashcan',
                   '413, assault_rifle', '414, backpack', '415, bakery', '416, balance_beam', '417, balloon',
                   '418, ballpoint', '419, Band_Aid', '420, banjo', '421, bannister', '422, barbell',
                   '423, barber_chair', '424, barbershop', '425, barn', '426, barometer', '427, barrel', '428, barrow',
                   '429, baseball', '430, basketball', '431, bassinet', '432, bassoon', '433, bathing_cap',
                   '434, bath_towel', '435, bathtub', '436, beach_wagon', '437, beacon', '438, beaker', '439, bearskin',
                   '440, beer_bottle', '441, beer_glass', '442, bell_cote', '443, bib', '444, bicycle-built-for-two',
                   '445, bikini', '446, binder', '447, binoculars', '448, birdhouse', '449, boathouse', '450, bobsled',
                   '451, bolo_tie', '452, bonnet', '453, bookcase', '454, bookshop', '455, bottlecap', '456, bow',
                   '457, bow_tie', '458, brass', '459, brassiere', '460, breakwater', '461, breastplate', '462, broom',
                   '463, bucket', '464, buckle', '465, bulletproof_vest', '466, bullet_train', '467, butcher_shop',
                   '468, cab', '469, caldron', '470, candle', '471, cannon', '472, canoe', '473, can_opener',
                   '474, cardigan', '475, car_mirror', '476, carousel', "477, carpenter's_kit", '478, carton',
                   '479, car_wheel', '480, cash_machine', '481, cassette', '482, cassette_player', '483, castle',
                   '484, catamaran', '485, CD_player', '486, cello', '487, cellular_telephone', '488, chain',
                   '489, chainlink_fence', '490, chain_mail', '491, chain_saw', '492, chest', '493, chiffonier',
                   '494, chime', '495, china_cabinet', '496, Christmas_stocking', '497, church', '498, cinema',
                   '499, cleaver', '500, cliff_dwelling', '501, cloak', '502, clog', '503, cocktail_shaker',
                   '504, coffee_mug', '505, coffeepot', '506, coil', '507, combination_lock', '508, computer_keyboard',
                   '509, confectionery', '510, container_ship', '511, convertible', '512, corkscrew', '513, cornet',
                   '514, cowboy_boot', '515, cowboy_hat', '516, cradle', '517, crane', '518, crash_helmet',
                   '519, crate', '520, crib', '521, Crock_Pot', '522, croquet_ball', '523, crutch', '524, cuirass',
                   '525, dam', '526, desk', '527, desktop_computer', '528, dial_telephone', '529, diaper',
                   '530, digital_clock', '531, digital_watch', '532, dining_table', '533, dishrag', '534, dishwasher',
                   '535, disk_brake', '536, dock', '537, dogsled', '538, dome', '539, doormat',
                   '540, drilling_platform', '541, drum', '542, drumstick', '543, dumbbell', '544, Dutch_oven',
                   '545, electric_fan', '546, electric_guitar', '547, electric_locomotive', '548, entertainment_center',
                   '549, envelope', '550, espresso_maker', '551, face_powder', '552, feather_boa', '553, file',
                   '554, fireboat', '555, fire_engine', '556, fire_screen', '557, flagpole', '558, flute',
                   '559, folding_chair', '560, football_helmet', '561, forklift', '562, fountain', '563, fountain_pen',
                   '564, four-poster', '565, freight_car', '566, French_horn', '567, frying_pan', '568, fur_coat',
                   '569, garbage_truck', '570, gasmask', '571, gas_pump', '572, goblet', '573, go-kart',
                   '574, golf_ball', '575, golfcart', '576, gondola', '577, gong', '578, gown', '579, grand_piano',
                   '580, greenhouse', '581, grille', '582, grocery_store', '583, guillotine', '584, hair_slide',
                   '585, hair_spray', '586, half_track', '587, hammer', '588, hamper', '589, hand_blower',
                   '590, hand-held_computer', '591, handkerchief', '592, hard_disc', '593, harmonica', '594, harp',
                   '595, harvester', '596, hatchet', '597, holster', '598, home_theater', '599, honeycomb', '600, hook',
                   '601, hoopskirt', '602, horizontal_bar', '603, horse_cart', '604, hourglass', '605, iPod',
                   '606, iron', "607, jack-o'-lantern", '608, jean', '609, jeep', '610, jersey', '611, jigsaw_puzzle',
                   '612, jinrikisha', '613, joystick', '614, kimono', '615, knee_pad', '616, knot', '617, lab_coat',
                   '618, ladle', '619, lampshade', '620, laptop', '621, lawn_mower', '622, lens_cap',
                   '623, letter_opener', '624, library', '625, lifeboat', '626, lighter', '627, limousine',
                   '628, liner', '629, lipstick', '630, Loafer', '631, lotion', '632, loudspeaker', '633, loupe',
                   '634, lumbermill', '635, magnetic_compass', '636, mailbag', '637, mailbox', '638, maillot',
                   '639, maillot', '640, manhole_cover', '641, maraca', '642, marimba', '643, mask', '644, matchstick',
                   '645, maypole', '646, maze', '647, measuring_cup', '648, medicine_chest', '649, megalith',
                   '650, microphone', '651, microwave', '652, military_uniform', '653, milk_can', '654, minibus',
                   '655, miniskirt', '656, minivan', '657, missile', '658, mitten', '659, mixing_bowl',
                   '660, mobile_home', '661, Model_T', '662, modem', '663, monastery', '664, monitor', '665, moped',
                   '666, mortar', '667, mortarboard', '668, mosque', '669, mosquito_net', '670, motor_scooter',
                   '671, mountain_bike', '672, mountain_tent', '673, mouse', '674, mousetrap', '675, moving_van',
                   '676, muzzle', '677, nail', '678, neck_brace', '679, necklace', '680, nipple', '681, notebook',
                   '682, obelisk', '683, oboe', '684, ocarina', '685, odometer', '686, oil_filter', '687, organ',
                   '688, oscilloscope', '689, overskirt', '690, oxcart', '691, oxygen_mask', '692, packet',
                   '693, paddle', '694, paddlewheel', '695, padlock', '696, paintbrush', '697, pajama', '698, palace',
                   '699, panpipe', '700, paper_towel', '701, parachute', '702, parallel_bars', '703, park_bench',
                   '704, parking_meter', '705, passenger_car', '706, patio', '707, pay-phone', '708, pedestal',
                   '709, pencil_box', '710, pencil_sharpener', '711, perfume', '712, Petri_dish', '713, photocopier',
                   '714, pick', '715, pickelhaube', '716, picket_fence', '717, pickup', '718, pier', '719, piggy_bank',
                   '720, pill_bottle', '721, pillow', '722, ping-pong_ball', '723, pinwheel', '724, pirate',
                   '725, pitcher', '726, plane', '727, planetarium', '728, plastic_bag', '729, plate_rack', '730, plow',
                   '731, plunger', '732, Polaroid_camera', '733, pole', '734, police_van', '735, poncho',
                   '736, pool_table', '737, pop_bottle', '738, pot', "739, potter's_wheel", '740, power_drill',
                   '741, prayer_rug', '742, printer', '743, prison', '744, projectile', '745, projector', '746, puck',
                   '747, punching_bag', '748, purse', '749, quill', '750, quilt', '751, racer', '752, racket',
                   '753, radiator', '754, radio', '755, radio_telescope', '756, rain_barrel',
                   '757, recreational_vehicle', '758, reel', '759, reflex_camera', '760, refrigerator',
                   '761, remote_control', '762, restaurant', '763, revolver', '764, rifle', '765, rocking_chair',
                   '766, rotisserie', '767, rubber_eraser', '768, rugby_ball', '769, rule', '770, running_shoe',
                   '771, safe', '772, safety_pin', '773, saltshaker', '774, sandal', '775, sarong', '776, sax',
                   '777, scabbard', '778, scale', '779, school_bus', '780, schooner', '781, scoreboard', '782, screen',
                   '783, screw', '784, screwdriver', '785, seat_belt', '786, sewing_machine', '787, shield',
                   '788, shoe_shop', '789, shoji', '790, shopping_basket', '791, shopping_cart', '792, shovel',
                   '793, shower_cap', '794, shower_curtain', '795, ski', '796, ski_mask', '797, sleeping_bag',
                   '798, slide_rule', '799, sliding_door', '800, slot', '801, snorkel', '802, snowmobile',
                   '803, snowplow', '804, soap_dispenser', '805, soccer_ball', '806, sock', '807, solar_dish',
                   '808, sombrero', '809, soup_bowl', '810, space_bar', '811, space_heater', '812, space_shuttle',
                   '813, spatula', '814, speedboat', '815, spider_web', '816, spindle', '817, sports_car',
                   '818, spotlight', '819, stage', '820, steam_locomotive', '821, steel_arch_bridge', '822, steel_drum',
                   '823, stethoscope', '824, stole', '825, stone_wall', '826, stopwatch', '827, stove', '828, strainer',
                   '829, streetcar', '830, stretcher', '831, studio_couch', '832, stupa', '833, submarine', '834, suit',
                   '835, sundial', '836, sunglass', '837, sunglasses', '838, sunscreen', '839, suspension_bridge',
                   '840, swab', '841, sweatshirt', '842, swimming_trunks', '843, swing', '844, switch', '845, syringe',
                   '846, table_lamp', '847, tank', '848, tape_player', '849, teapot', '850, teddy', '851, television',
                   '852, tennis_ball', '853, thatch', '854, theater_curtain', '855, thimble', '856, thresher',
                   '857, throne', '858, tile_roof', '859, toaster', '860, tobacco_shop', '861, toilet_seat',
                   '862, torch', '863, totem_pole', '864, tow_truck', '865, toyshop', '866, tractor',
                   '867, trailer_truck', '868, tray', '869, trench_coat', '870, tricycle', '871, trimaran',
                   '872, tripod', '873, triumphal_arch', '874, trolleybus', '875, trombone', '876, tub',
                   '877, turnstile', '878, typewriter_keyboard', '879, umbrella', '880, unicycle', '881, upright',
                   '882, vacuum', '883, vase', '884, vault', '885, velvet', '886, vending_machine', '887, vestment',
                   '888, viaduct', '889, violin', '890, volleyball', '891, waffle_iron', '892, wall_clock',
                   '893, wallet', '894, wardrobe', '895, warplane', '896, washbasin', '897, washer',
                   '898, water_bottle', '899, water_jug', '900, water_tower', '901, whiskey_jug', '902, whistle',
                   '903, wig', '904, window_screen', '905, window_shade', '906, Windsor_tie', '907, wine_bottle',
                   '908, wing', '909, wok', '910, wooden_spoon', '911, wool', '912, worm_fence', '913, wreck',
                   '914, yawl', '915, yurt', '916, web_site', '917, comic_book', '918, crossword_puzzle',
                   '919, street_sign', '920, traffic_light', '921, book_jacket', '922, menu', '923, plate',
                   '924, guacamole', '925, consomme', '926, hot_pot', '927, trifle', '928, ice_cream', '929, ice_lolly',
                   '930, French_loaf', '931, bagel', '932, pretzel', '933, cheeseburger', '934, hotdog',
                   '935, mashed_potato', '936, head_cabbage', '937, broccoli', '938, cauliflower', '939, zucchini',
                   '940, spaghetti_squash', '941, acorn_squash', '942, butternut_squash', '943, cucumber',
                   '944, artichoke', '945, bell_pepper', '946, cardoon', '947, mushroom', '948, Granny_Smith',
                   '949, strawberry', '950, orange', '951, lemon', '952, fig', '953, pineapple', '954, banana',
                   '955, jackfruit', '956, custard_apple', '957, pomegranate', '958, hay', '959, carbonara',
                   '960, chocolate_sauce', '961, dough', '962, meat_loaf', '963, pizza', '964, potpie', '965, burrito',
                   '966, red_wine', '967, espresso', '968, cup', '969, eggnog', '970, alp', '971, bubble', '972, cliff',
                   '973, coral_reef', '974, geyser', '975, lakeside', '976, promontory', '977, sandbar',
                   '978, seashore', '979, valley', '980, volcano', '981, ballplayer', '982, groom', '983, scuba_diver',
                   '984, rapeseed', '985, daisy', "986, yellow_lady's_slipper", '987, corn', '988, acorn', '989, hip',
                   '990, buckeye', '991, coral_fungus', '992, agaric', '993, gyromitra', '994, stinkhorn',
                   '995, earthstar', '996, hen-of-the-woods', '997, bolete', '998, ear', '999, toilet_tissue']
        return classes

    def classify_image(self, image):
        # 对输入图片进行预处理
        transform = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(
                mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225]
            )
        ])
        image_tensor = transform(image).unsqueeze(0)
        # 使用模型进行预测
        with torch.no_grad():
            output = self.model(image_tensor)
            _, predicted = torch.max(output.data, 1)
            score = torch.nn.functional.softmax(output, dim=1)[0][predicted[0]].item()

        image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
        # 在图片上绘制预测结果
        label = f"{self.classes[predicted[0]]} ({score:.2f})"
        cv2.putText(image, label, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2)
        image = cv2.cvtColor(np.array(image), cv2.COLOR_BGR2RGB)
        return image, self.classes[predicted[0]]

    def classify_images(self, images):
        results = []
        for image in images:
            result = self.classify_image(image)
            results.append(result)
        return results

    def classify_interactive(self):
        examples = [
            ['bus.jpg', 'VGG16'],
            ['dogs.jpg', 'VGG16'],
            ['zidane.jpg', 'VGG16']
        ]
        input_image = gr.inputs.Image(type='pil')
        output_image = gr.outputs.Image(type='pil')
        output_text = gr.outputs.Textbox()
        model_name = gr.inputs.Dropdown(
            ["ResNet18", "AlexNet", "VGG16", "GoogLeNet", "ResNet50", "DenseNet121", "MobileNetV2"], label="Model",
            default="VGG16")

        def predict(image, model_name):
            self.model = self.load_model(model_name)
            result_image, result_class = self.classify_image(image)
            return result_image, result_class

        gr.Interface(fn=predict, inputs=[input_image, model_name], outputs=[output_image, output_text],
                     examples=examples, live=True, capture_session=True,
                     title="Image Classification", description="Upload an image and classify it.").launch()


if __name__ == "__main__":
    classifier = ImageClassifier("VGG16")
    classifier.classify_interactive()