Spaces:
Sleeping
Sleeping
import os | |
os.system("pip install timm~=0.9.2") | |
import gradio as gr | |
import timm | |
import torch | |
import torchvision.transforms as T | |
from PIL import Image, ImageDraw, ImageFont | |
import warnings | |
warnings.filterwarnings("ignore") | |
# 获取 ImageNet 类别列表 | |
imagenet_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'] | |
def download_test_img(): | |
# 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 preprocess(image: Image.Image) -> torch.Tensor: | |
transform = T.Compose([ | |
T.Resize(256), | |
T.CenterCrop(224), | |
T.ToTensor(), | |
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), | |
]) | |
return transform(image).unsqueeze(0) | |
# 后处理 | |
def postprocess(output: torch.Tensor) -> (str, float): | |
probabilities = torch.nn.functional.softmax(output, dim=1)[0] * 100 | |
_, index = output.max(1) | |
label = imagenet_classes[index.item()] | |
confidence = probabilities[index].item() | |
return label, confidence | |
def predict(image: Image.Image, model_name: str) -> (str, float): | |
model = get_model(model_name) | |
input_tensor = preprocess(image) | |
with torch.no_grad(): | |
output = model(input_tensor) | |
return postprocess(output) | |
def get_model(model_name: str): | |
try: | |
model = timm.create_model(model_name, pretrained=True) | |
model.eval() | |
except ValueError: | |
raise ValueError(f"Model {model_name} not found or pretrained weights not available.") | |
return model | |
# 将分类名称和置信度绘制在图像上 | |
def draw_label_on_image(image: Image.Image, label: str, confidence: float) -> Image.Image: | |
draw = ImageDraw.Draw(image) | |
font = ImageFont.load_default() | |
text = f"{label} ({confidence:.2f}%)" | |
draw.text((10, 10), text, font=font, fill=(255, 255, 255, 128)) | |
return image | |
def predict_and_draw(image: Image.Image, model_name: str) -> Image.Image: | |
label, confidence = predict(image, model_name) | |
image_output = draw_label_on_image(image, label, confidence) | |
return image_output, label, confidence | |
examples = [ | |
['bus.jpg', 'resnet50'], | |
['dogs.jpg', 'efficientnet_b0'], | |
['zidane.jpg', 'vgg16_bn'] | |
] | |
download_test_img() | |
# 定义输入和输出 | |
image_input = gr.inputs.Image(type="pil") | |
model_input = gr.inputs.Dropdown( | |
choices=["resnet50", "efficientnet_b0", "efficientnet_b1", "efficientnet_b2", "efficientnet_b3", | |
"efficientnet_b4", "efficientnet_b5", "efficientnet_b6", "efficientnet_b7", "mobilenetv3_large_100", | |
"mobilenetv3_rw", "resnet18", "resnet34", "resnet101", "resnext50_32x4d", "resnext101_32x8d", | |
"vgg11_bn", "vgg13_bn", "vgg16_bn", "vgg19_bn"], label="Model", default="resnet50") | |
image_output = gr.outputs.Image(type="pil") | |
text_output = gr.outputs.Textbox(label="Label", type="text") | |
score_output = gr.outputs.Textbox(label="Confidence", type="text") | |
# 创建 Gradio 界面 | |
iface = gr.Interface(fn=predict_and_draw, inputs=[image_input, model_input], | |
outputs=[image_output, text_output, score_output], examples=examples, | |
title="Image Classification on timm", description="Upload an image and classify,it.") | |
# 启动 Gradio 界面 | |
iface.launch() | |