Spaces:
Runtime error
Runtime error
# Check Pytorch installation | |
import torch, torchvision | |
print("torch version:",torch.__version__, "cuda:",torch.cuda.is_available()) | |
# Check MMDetection installation | |
import mmdet | |
import os | |
import mmcv | |
import mmengine | |
from mmdet.apis import init_detector, inference_detector | |
from mmdet.utils import register_all_modules | |
from mmdet.registry import VISUALIZERS | |
from huggingface_hub import hf_hub_download | |
from huggingface_hub import snapshot_download | |
from time import time | |
classes = ['Beach', | |
'Sea', | |
'Wave', | |
'Rock', | |
'Breaking wave', | |
'Reflection of the sea', | |
'Foam', | |
'Algae', | |
'Vegetation', | |
'Watermark', | |
'Bird', | |
'Ship', | |
'Boat', | |
'Car', | |
'Kayak', | |
"Shark's line", | |
'Dock', | |
'Dog', | |
'Unidentifiable shade', | |
'Bird shadow', | |
'Boat shadow', | |
'Kayal shade', | |
'Surfer shadow', | |
'Shark shadow', | |
'Surfboard shadow', | |
'Crocodile', | |
'Sea cow', | |
'Stingray', | |
'Person', | |
'ocean', | |
'Surfer', | |
'Surfer', | |
'Fish', | |
'Killer whale', | |
'Whale', | |
'Dolphin', | |
'Miscellaneous', | |
'Unidentifiable shark', | |
'Carpet shark', | |
'Dusty shark', | |
'Blue shark', | |
'Great white shark', | |
'Copper shark', | |
'Nurse shark', | |
'Silky shark', | |
'Leopard shark', | |
'Shortfin mako shark', | |
'Hammerhead shark', | |
'Oceanic whitetip shark', | |
'Blacktip shark', | |
'Tiger shark', | |
'Bull shark']*3 | |
class_sizes = {'Beach': None, | |
'Sea': None, | |
'Wave': None, | |
'Rock': None, | |
'Breaking wave': None, | |
'Reflection of the sea': None, | |
'Foam': None, | |
'Algae': None, | |
'Vegetation': None, | |
'Watermark': None, | |
'Bird': {'feet':[1, 3], 'meter': [0.3, 0.9], 'kg': [0.5, 1.5], 'pounds': [1, 3]}, | |
'Ship': {'feet':[10, 100], 'meter': [3, 30], 'kg': [1000, 100000], 'pounds': [2200, 220000]}, | |
'Boat': {'feet':[10, 45], 'meter': [3, 15], 'kg': [750, 80000], 'pounds': [1500, 160000]}, | |
'Car': {'feet':[10, 20], 'meter': [3, 6], 'kg': [1000, 2000], 'pounds': [2200, 4400]}, | |
'Kayak': {'feet':[10, 20], 'meter': [3, 6], 'kg': [50, 300], 'pounds': [100, 600]}, | |
"Shark's line": None, | |
'Dock': None, | |
'Dog': {'feet':[1, 3], 'meter': [0.3, 0.9], 'kg': [10, 50], 'pounds': [20, 100]}, | |
'Unidentifiable shade': None, | |
'Bird shadow': None, | |
'Boat shadow': None, | |
'Kayal shade': None, | |
'Surfer shadow': None, | |
'Shark shadow': None, | |
'Surfboard shadow': None, | |
'Crocodile': {'feet':[10, 20], 'meter': [3, 6], 'kg': [410, 1000], 'pounds': [900, 2200]}, | |
'Sea cow': {'feet':[9,12], 'meter': [3, 4], 'kg': [400, 590], 'pounds': [900, 1300]}, | |
'Stingray': {'feet':[2, 7.5], 'meter': [0.6, 2.5], 'kg': [100, 300], 'pounds': [220, 770]}, | |
'Person': {'feet':[5, 7], 'meter': [1.5, 2.1], 'kg': [50, 150], 'pounds': [110, 300]}, | |
'Ocean': None, | |
'Surfer': {'feet':[5, 7], 'meter': [1.5, 2.1], 'kg': [50, 150], 'pounds': [110, 300]}, | |
'Surfer': {'feet':[5, 7], 'meter': [1.5, 2.1], 'kg': [50, 150], 'pounds': [110, 300]}, | |
'Fish': {'feet':[1, 3], 'meter': [0.3, 0.9], 'kg': [20, 150], 'pounds': [40, 300]}, | |
'Killer whale': {'feet':[10, 20], 'meter': [3, 6], 'kg': [3600, 5400], 'pounds': [8000, 12000]}, | |
'Whale': {'feet':[15, 30], 'meter': [4.5, 10], 'kg': [2500, 80000], 'pounds': [55000, 176000]}, | |
'Dolphin': {'feet':[6.6, 13.1], 'meter': [2, 4], 'kg': [150, 650], 'pounds': [330, 1430]}, | |
'Miscellaneous': None, | |
'Unidentifiable shark': {'feet': [2, 15], 'meter': [0.6, 4.5], 'kg': [50, 1000], 'pounds': [110, 2200]}, | |
'Carpet shark': {'feet': [4, 10], 'meter': [1.25, 3], 'kg': [50, 1000], 'pounds': [110, 2200]}, # Prob incorrect | |
'Dusty shark': {'feet': [9, 14], 'meter': [3, 4.25], 'kg': [160, 180], 'pounds': [350, 400]}, | |
'Blue shark': {'feet': [7.9, 12.5], 'meter': [2.4, 3], 'kg': [60, 120], 'pounds': [130, 260]}, | |
'Great white shark': {'feet': [13.1, 20], 'meter': [4, 6], 'kg': [680, 1800], 'pounds': [1500, 4000]}, | |
'Copper shark': {'feet': [7.2, 10.8], 'meter': [2.2, 3.3], 'kg': [130, 300], 'pounds': [290, 660]}, | |
'Nurse shark': {'feet': [7.9, 9.8], 'meter': [2.4, 3], 'kg': [90, 115], 'pounds': [200, 250]}, | |
'Silky shark': {'feet': [6.6, 8.2], 'meter': [2, 2.5], 'kg': [300, 380], 'pounds': [660, 840]}, | |
'Leopard shark': {'feet': [3.9, 4.9], 'meter': [1.2, 1.5], 'kg': [11, 20], 'pounds': [22, 44]}, | |
'Shortfin mako shark': {'feet': [10.5, 12], 'meter': [3.2, 3.6], 'kg': [60, 135], 'pounds': [130, 300]}, | |
'Hammerhead shark': {'feet': [4.9, 20], 'meter': [1.5, 6.1], 'kg': [230, 450], 'pounds': [500, 1000]}, | |
'Oceanic whitetip shark': {'feet': [5.9, 9.8], 'meter': [1.8, 3], 'kg': [36, 170], 'pounds': [80, 375]}, | |
'Blacktip shark': {'feet': [4.9, 6.6], 'meter': [1.5, 2], 'kg': [40, 100], 'pounds': [90, 220]}, | |
'Tiger shark': {'feet': [9.8, 18], 'meter': [3, 5.5], 'kg': [385, 635], 'pounds': [850, 1400]}, | |
'Bull shark': {'feet': [7.9, 11.2], 'meter': [2.4, 3.4], 'kg': [200, 315], 'pounds': [440, 690]}, | |
} | |
class_sizes_lower = {k.lower(): v for k, v in class_sizes.items()} | |
classes_is_shark = [1 if 'shark' in x.lower() else 0 for x in classes] | |
classes_is_human = [1 if 'person' or 'surfer' in x.lower() else 0 for x in classes] | |
classes_is_unknown = [1 if 'unidentifiable' in x.lower() else 0 for x in classes] | |
classes_is_shark_id = [i for i, x in enumerate(classes_is_shark) if x == 1] | |
classes_is_human_id = [i for i, x in enumerate(classes_is_human) if x == 1] | |
classes_is_unknown_id = [i for i, x in enumerate(classes_is_unknown) if x == 1] | |
REPO_ID = "SharkSpace/maskformer_model" | |
FILENAME = "mask2former" | |
snapshot_download(repo_id=REPO_ID, token= os.environ.get('SHARK_MODEL'),local_dir='model/') | |
# Choose to use a config and initialize the detector | |
config_file ='model/mask2former_swin-t-p4-w7-224_8xb2-lsj-50e_coco-panoptic/mask2former_swin-t-p4-w7-224_8xb2-lsj-50e_coco-panoptic.py' | |
#'/content/mmdetection/configs/panoptic_fpn/panoptic-fpn_r50_fpn_ms-3x_coco.py' | |
# Setup a checkpoint file to load | |
checkpoint_file ='model/mask2former_swin-t-p4-w7-224_8xb2-lsj-50e_coco-panoptic/checkpoint_v2.pth' | |
# '/content/drive/MyDrive/Algorithms/weights/shark_panoptic_weights_16_4_23/panoptic-fpn_r50_fpn_ms-3x_coco/epoch_36.pth' | |
# register all modules in mmdet into the registries | |
register_all_modules() | |
# build the model from a config file and a checkpoint file | |
model = init_detector(config_file, checkpoint_file, device='cuda:0') # or device='cuda:0' | |
model.dataset_meta['palette'] = model.dataset_meta['palette'] + model.dataset_meta['palette'][-23:] | |
model.dataset_meta['classes'] = classes | |
print(model.cfg.visualizer) | |
# init visualizer(run the block only once in jupyter notebook) | |
visualizer = VISUALIZERS.build(model.cfg.visualizer) | |
print(dir(visualizer)) | |
# the dataset_meta is loaded from the checkpoint and | |
# then pass to the model in init_detector | |
visualizer.dataset_meta = model.dataset_meta | |
classes = visualizer.dataset_meta.get('classes', None) | |
palette = visualizer.dataset_meta.get('palette', None) | |
print(len(classes)) | |
print(len(palette)) | |
def inference_frame_serial(image, visualize = True): | |
#start = time() | |
result = inference_detector(model, image) | |
#print(f'inference time: {time()-start}') | |
# show the results | |
if visualize: | |
visualizer.add_datasample( | |
'result', | |
image, | |
data_sample=result, | |
draw_gt = None, | |
show=False | |
) | |
frame = visualizer.get_image() | |
else: | |
frame = None | |
return frame, result | |
def inference_frame(image): | |
result = inference_detector(model, image) | |
# show the results | |
frames = [] | |
cnt=0 | |
for res in result: | |
visualizer.add_datasample( | |
'result', | |
image[cnt], | |
data_sample=res.numpy(), | |
draw_gt = None, | |
show=False | |
) | |
frame = visualizer.get_image() | |
frames.append(frame) | |
cnt+=1 | |
#frames = process_frames(result, image, visualizer) | |
return frames | |
def inference_frame_par_ready(image): | |
result = inference_detector(model, image) | |
return [result[i].numpy() for i in range(len(result))] | |
def process_frame(in_tuple = (None, None, None)): | |
visualizer.add_datasample( | |
'result', | |
in_tuple[1], #image, | |
data_sample=in_tuple[0], #result | |
draw_gt = None, | |
show=False | |
) | |
#frame = visualizer.get_image() | |
#print(in_tuple[2]) | |
return visualizer.get_image() | |
#def process_frame(frame): | |
# def process_frames(result, image, visualizer): | |
# frames = [] | |
# lock = threading.Lock() | |
# def process_data(cnt, res, img): | |
# visualizer.add_datasample('result', img, data_sample=res, draw_gt=None, show=False) | |
# frame = visualizer.get_image() | |
# with lock: | |
# frames.append(frame) | |
# threads = [] | |
# for cnt, res in enumerate(result): | |
# t = threading.Thread(target=process_data, args=(cnt, res, image[cnt])) | |
# threads.append(t) | |
# t.start() | |
# for t in threads: | |
# t.join() | |
# return frames | |