import os
os.environ['CUDA_VISIBLE_DEVICES'] = '2'
import torch
torch.rand(10)
import torch.nn as nn
import torch.nn.functional as F
import glob
from tqdm import tqdm, trange
print(torch.cuda.is_available())
print(torch.cuda.get_device_name())
print(torch.cuda.current_device())

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('Using device:', device)
print()

#Additional Info when using cuda
if device.type == 'cuda':
    print(torch.cuda.get_device_name(0))
    print('Memory Usage:')
    print('Allocated:', round(torch.cuda.memory_allocated(0)/1024**3,1), 'GB')
    print('Cached:   ', round(torch.cuda.memory_reserved(0)/1024**3,1), 'GB')
    
import torch.backends.cudnn as cudnn
import numpy as np
import os, cv2
from tqdm import tqdm, trange
import seaborn as sns
from models.experimental import attempt_load
from utils.datasets import LoadStreams, LoadImages
from utils.general import (
    check_img_size, non_max_suppression, apply_classifier, scale_coords, xyxy2xywh, plot_one_box, strip_optimizer)
from utils.torch_utils import select_device, load_classifier, time_synchronized
from my_utils import xyxy_2_xyxyo, draw_boxes
# Initialize
device = select_device('')
half = device.type != 'cpu'  # half precision only supported on CUDA

def prepare_input(img1, img_size=416, half=True):
    
    img2 = cv2.resize(img1, (img_size, img_size)) # W x H
    img2 = img2.transpose(2,0,1)
    img2 = img2[np.newaxis, ...]
    img2 = torch.from_numpy(img2).to(device) # torch image is ch x H x W
    img2 = img2.half() if not half else img2.float()
    img2 /= 255.0
    
    return img2
#%%

# Directories
out = '/home/user01/data_ssd/Talha/yolo/op/'
weights = '/home/user01/data_ssd/Talha/yolo/ScaledYOLOv4/runs/exp2_yolov4-csp-results/weights/best_yolov4-csp-results.pt'
source = '/home/user01/data_ssd/Talha/yolo/paprika_y5/valid/images/'
imgsz = 416
conf_thres = 0.4
iou_thres = 0.5
classes =  [0,1,2,3,4,5] 
class_names = ["blossom_end_rot", "graymold","powdery_mildew","spider_mite",
               "spotting_disease", "snails_and_slugs"] 

# deleting files in op_dir
filelist = [ f for f in os.listdir(out)]# if f.endswith(".png") ]
for f in tqdm(filelist, desc = 'Deleting old files fro directory'):
    os.remove(os.path.join(out, f))

# Load model
model = attempt_load(weights, map_location=device)  # load FP32 model
imgsz = check_img_size(imgsz, s=model.stride.max())  # check img_size

if half:
    model.half()  # to FP16

# Load model
model = attempt_load(weights, map_location=device)  # load FP32 model
imgsz = check_img_size(imgsz, s=model.stride.max())  # check img_size

img_paths = glob.glob('/home/user01/data_ssd/Talha/yolo/paprika_y5/test/images/*.png') + \
            glob.glob('/home/user01/data_ssd/Talha/yolo/paprika_y5/test/images/*.jpg')

# Run inference
if device.type != 'cpu':
    model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters())))  # run once
#%%
for i in trange(len(img_paths)):
    
    
    path = img_paths[i]
    img1  = cv2.imread(path)
    img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2RGB)
    img_h, img_w, _ = img1.shape
    
    img2 = prepare_input(img1, 416, half)
    # get file name
    name = os.path.basename(path)[:-4]
    # Inference
    t1 = time_synchronized()
    pred = model(img2, augment=False)[0]

    # Apply NMS
    pred = non_max_suppression(pred, conf_thres, iou_thres, classes=classes, agnostic=True)
    if pred[0] is not None:
        boxes = pred[0].cpu().detach().numpy() # <xmin><ymin><xmax><ymax><confd><class_id>
    else:
        boxes = np.array([10.0, 20.0, 30.0, 50.0, 0.75, 0]).reshape(1,6) # dummy values
    
    coords_minmax = np.zeros((boxes.shape[0], 4)) # droping 5th value
    confd = np.zeros((boxes.shape[0], 1))
    class_ids = np.zeros((boxes.shape[0], 1))
    # assign
    coords_minmax = boxes[:,0:4] # coords
    confd = boxes[:,4] # confidence
    class_ids = boxes[:,5] # class id
    
    coords_xyminmax = []
    det_classes = []
    for i in range(boxes.shape[0]):
        coords_xyminmax.append(xyxy_2_xyxyo(img_w, img_h, coords_minmax[i]))
        det_classes.append(class_names[int(class_ids[i])])
    all_bounding_boxnind = []
    for i in range(boxes.shape[0]):
        
        bounding_box = [0.0] * 6
        
        bounding_box[0] = det_classes[i]
        bounding_box[1] = confd[i]
        bounding_box[2] = coords_xyminmax[i][0]
        bounding_box[3] = coords_xyminmax[i][1]
        bounding_box[4] = coords_xyminmax[i][2]
        bounding_box[5] = coords_xyminmax[i][3]
        
        bounding_box = str(bounding_box)[1:-1]# remove square brackets
        bounding_box = bounding_box.replace("'",'')# removing inverted commas around class name
        bounding_box = "".join(bounding_box.split())# remove spaces in between **here dont give space inbetween the inverted commas "".
        all_bounding_boxnind.append(bounding_box)
    
    all_bounding_boxnind = ' '.join(map(str, all_bounding_boxnind))# convert list to string
    all_bounding_boxnind=list(all_bounding_boxnind.split(' ')) # convert strin to list
    # replacing commas with spaces
    for i in range(len(all_bounding_boxnind)):
        all_bounding_boxnind[i] = all_bounding_boxnind[i].replace(',',' ')
    for i in range(len(all_bounding_boxnind)):
    # check if file exiscts else make new
        with open(out +'{}.txt'.format(name), "a+") as file_object:
            # Move read cursor to the start of file.
            file_object.seek(0)
            # If file is not empty then append '\n'
            data = file_object.read(100)
            if len(data) > 0 :
                file_object.write("\n")
            # Append text at the end of file
            file_object.write(all_bounding_boxnind[i])
#%%
import glob, random
import matplotlib.pyplot as plt
import matplotlib as mpl
mpl.rcParams['figure.dpi'] = 300

img_paths = glob.glob('/home/user01/data_ssd/Talha/yolo/paprika_y5/test/images/*.png') + \
            glob.glob('/home/user01/data_ssd/Talha/yolo/paprika_y5/test/images/*.jpg')
img_path = random.choice(img_paths)

img1  = cv2.imread(img_path)
img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2RGB)
img_h, img_w, _ = img1.shape

img2 = prepare_input(img1, 416, half)

pred = model(img2, augment=False)[0]

# Apply NMS
pred = non_max_suppression(pred, conf_thres, iou_thres, classes=classes, agnostic=True)
boxes = pred[0].cpu().detach().numpy() # <xmin><ymin><xmax><ymax><confd><class_id>

coords_minmax = np.zeros((boxes.shape[0], 4)) # droping 5th value
confd = np.zeros((boxes.shape[0], 1))
class_ids = np.zeros((boxes.shape[0], 1))
# assign
coords_minmax = boxes[:,0:4] # coords
confd = boxes[:,4] # confidence
class_ids = boxes[:,5] # class id
    
coords_xyminmax = []
det_classes = []
for i in range(boxes.shape[0]):
    coords_xyminmax.append(xyxy_2_xyxyo(img_w, img_h, coords_minmax[i]))
    det_classes.append(class_names[int(class_ids[i])])

t = np.asarray(coords_xyminmax)
op = draw_boxes(img1, confd, t, det_classes, class_names, order='xy_minmax', analysis=False)
plt.imshow(op)
print('='*50)
print('Image Name: ', os.path.basename(img_path),img1.shape)
print('\nClass_name      ', '| B_box Coords ', '| Confidence')
print('_'*50)
for k in range(len(det_classes)):
    print(det_classes[k], t[k], confd[k])
print('='*50)
[((64, 78), 'torch.rand', 'torch.rand', (['(10)'], {}), '(10)\n', (74, 78), False, 'import torch\n'), ((1189, 1206), 'utils.torch_utils.select_device', 'select_device', (['""""""'], {}), "('')\n", (1202, 1206), False, 'from utils.torch_utils import select_device, load_classifier, time_synchronized\n'), ((2183, 2238), 'tqdm.tqdm', 'tqdm', (['filelist'], {'desc': '"""Deleting old files fro directory"""'}), "(filelist, desc='Deleting old files fro directory')\n", (2187, 2238), False, 'from tqdm import tqdm, trange\n'), ((2300, 2342), 'models.experimental.attempt_load', 'attempt_load', (['weights'], {'map_location': 'device'}), '(weights, map_location=device)\n', (2312, 2342), False, 'from models.experimental import attempt_load\n'), ((2492, 2534), 'models.experimental.attempt_load', 'attempt_load', (['weights'], {'map_location': 'device'}), '(weights, map_location=device)\n', (2504, 2534), False, 'from models.experimental import attempt_load\n'), ((6058, 6082), 'random.choice', 'random.choice', (['img_paths'], {}), '(img_paths)\n', (6071, 6082), False, 'import glob, random\n'), ((6092, 6112), 'cv2.imread', 'cv2.imread', (['img_path'], {}), '(img_path)\n', (6102, 6112), False, 'import os, cv2\n'), ((6120, 6157), 'cv2.cvtColor', 'cv2.cvtColor', (['img1', 'cv2.COLOR_BGR2RGB'], {}), '(img1, cv2.COLOR_BGR2RGB)\n', (6132, 6157), False, 'import os, cv2\n'), ((6284, 6369), 'utils.general.non_max_suppression', 'non_max_suppression', (['pred', 'conf_thres', 'iou_thres'], {'classes': 'classes', 'agnostic': '(True)'}), '(pred, conf_thres, iou_thres, classes=classes, agnostic=True\n    )\n', (6303, 6369), False, 'from utils.general import check_img_size, non_max_suppression, apply_classifier, scale_coords, xyxy2xywh, plot_one_box, strip_optimizer\n'), ((6465, 6494), 'numpy.zeros', 'np.zeros', (['(boxes.shape[0], 4)'], {}), '((boxes.shape[0], 4))\n', (6473, 6494), True, 'import numpy as np\n'), ((6523, 6552), 'numpy.zeros', 'np.zeros', (['(boxes.shape[0], 1)'], {}), '((boxes.shape[0], 1))\n', (6531, 6552), True, 'import numpy as np\n'), ((6565, 6594), 'numpy.zeros', 'np.zeros', (['(boxes.shape[0], 1)'], {}), '((boxes.shape[0], 1))\n', (6573, 6594), True, 'import numpy as np\n'), ((6916, 6943), 'numpy.asarray', 'np.asarray', (['coords_xyminmax'], {}), '(coords_xyminmax)\n', (6926, 6943), True, 'import numpy as np\n'), ((6949, 7040), 'my_utils.draw_boxes', 'draw_boxes', (['img1', 'confd', 't', 'det_classes', 'class_names'], {'order': '"""xy_minmax"""', 'analysis': '(False)'}), "(img1, confd, t, det_classes, class_names, order='xy_minmax',\n    analysis=False)\n", (6959, 7040), False, 'from my_utils import xyxy_2_xyxyo, draw_boxes\n'), ((7037, 7051), 'matplotlib.pyplot.imshow', 'plt.imshow', (['op'], {}), '(op)\n', (7047, 7051), True, 'import matplotlib.pyplot as plt\n'), ((181, 206), 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), '()\n', (204, 206), False, 'import torch\n'), ((214, 242), 'torch.cuda.get_device_name', 'torch.cuda.get_device_name', ([], {}), '()\n', (240, 242), False, 'import torch\n'), ((250, 277), 'torch.cuda.current_device', 'torch.cuda.current_device', ([], {}), '()\n', (275, 277), False, 'import torch\n'), ((1343, 1381), 'cv2.resize', 'cv2.resize', (['img1', '(img_size, img_size)'], {}), '(img1, (img_size, img_size))\n', (1353, 1381), False, 'import os, cv2\n'), ((2637, 2711), 'glob.glob', 'glob.glob', (['"""/home/user01/data_ssd/Talha/yolo/paprika_y5/test/images/*.png"""'], {}), "('/home/user01/data_ssd/Talha/yolo/paprika_y5/test/images/*.png')\n", (2646, 2711), False, 'import glob, random\n'), ((2728, 2802), 'glob.glob', 'glob.glob', (['"""/home/user01/data_ssd/Talha/yolo/paprika_y5/test/images/*.jpg"""'], {}), "('/home/user01/data_ssd/Talha/yolo/paprika_y5/test/images/*.jpg')\n", (2737, 2802), False, 'import glob, random\n'), ((3028, 3044), 'cv2.imread', 'cv2.imread', (['path'], {}), '(path)\n', (3038, 3044), False, 'import os, cv2\n'), ((3056, 3093), 'cv2.cvtColor', 'cv2.cvtColor', (['img1', 'cv2.COLOR_BGR2RGB'], {}), '(img1, cv2.COLOR_BGR2RGB)\n', (3068, 3093), False, 'import os, cv2\n'), ((3258, 3277), 'utils.torch_utils.time_synchronized', 'time_synchronized', ([], {}), '()\n', (3275, 3277), False, 'from utils.torch_utils import select_device, load_classifier, time_synchronized\n'), ((3347, 3432), 'utils.general.non_max_suppression', 'non_max_suppression', (['pred', 'conf_thres', 'iou_thres'], {'classes': 'classes', 'agnostic': '(True)'}), '(pred, conf_thres, iou_thres, classes=classes, agnostic=True\n    )\n', (3366, 3432), False, 'from utils.general import check_img_size, non_max_suppression, apply_classifier, scale_coords, xyxy2xywh, plot_one_box, strip_optimizer\n'), ((3670, 3699), 'numpy.zeros', 'np.zeros', (['(boxes.shape[0], 4)'], {}), '((boxes.shape[0], 4))\n', (3678, 3699), True, 'import numpy as np\n'), ((3732, 3761), 'numpy.zeros', 'np.zeros', (['(boxes.shape[0], 1)'], {}), '((boxes.shape[0], 1))\n', (3740, 3761), True, 'import numpy as np\n'), ((3778, 3807), 'numpy.zeros', 'np.zeros', (['(boxes.shape[0], 1)'], {}), '((boxes.shape[0], 1))\n', (3786, 3807), True, 'import numpy as np\n'), ((5881, 5955), 'glob.glob', 'glob.glob', (['"""/home/user01/data_ssd/Talha/yolo/paprika_y5/test/images/*.png"""'], {}), "('/home/user01/data_ssd/Talha/yolo/paprika_y5/test/images/*.png')\n", (5890, 5955), False, 'import glob, random\n'), ((5972, 6046), 'glob.glob', 'glob.glob', (['"""/home/user01/data_ssd/Talha/yolo/paprika_y5/test/images/*.jpg"""'], {}), "('/home/user01/data_ssd/Talha/yolo/paprika_y5/test/images/*.jpg')\n", (5981, 6046), False, 'import glob, random\n'), ((7088, 7114), 'os.path.basename', 'os.path.basename', (['img_path'], {}), '(img_path)\n', (7104, 7114), False, 'import os, cv2\n'), ((312, 337), 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), '()\n', (335, 337), False, 'import torch\n'), ((459, 488), 'torch.cuda.get_device_name', 'torch.cuda.get_device_name', (['(0)'], {}), '(0)\n', (485, 488), False, 'import torch\n'), ((2132, 2147), 'os.listdir', 'os.listdir', (['out'], {}), '(out)\n', (2142, 2147), False, 'import os, cv2\n'), ((2256, 2276), 'os.path.join', 'os.path.join', (['out', 'f'], {}), '(out, f)\n', (2268, 2276), False, 'import os, cv2\n'), ((3205, 3227), 'os.path.basename', 'os.path.basename', (['path'], {}), '(path)\n', (3221, 3227), False, 'import os, cv2\n'), ((6810, 6854), 'my_utils.xyxy_2_xyxyo', 'xyxy_2_xyxyo', (['img_w', 'img_h', 'coords_minmax[i]'], {}), '(img_w, img_h, coords_minmax[i])\n', (6822, 6854), False, 'from my_utils import xyxy_2_xyxyo, draw_boxes\n'), ((1467, 1489), 'torch.from_numpy', 'torch.from_numpy', (['img2'], {}), '(img2)\n', (1483, 1489), False, 'import torch\n'), ((4055, 4099), 'my_utils.xyxy_2_xyxyo', 'xyxy_2_xyxyo', (['img_w', 'img_h', 'coords_minmax[i]'], {}), '(img_w, img_h, coords_minmax[i])\n', (4067, 4099), False, 'from my_utils import xyxy_2_xyxyo, draw_boxes\n'), ((547, 577), 'torch.cuda.memory_allocated', 'torch.cuda.memory_allocated', (['(0)'], {}), '(0)\n', (574, 577), False, 'import torch\n'), ((626, 655), 'torch.cuda.memory_reserved', 'torch.cuda.memory_reserved', (['(0)'], {}), '(0)\n', (652, 655), False, 'import torch\n'), ((3573, 3616), 'numpy.array', 'np.array', (['[10.0, 20.0, 30.0, 50.0, 0.75, 0]'], {}), '([10.0, 20.0, 30.0, 50.0, 0.75, 0])\n', (3581, 3616), True, 'import numpy as np\n'), ((2855, 2886), 'torch.zeros', 'torch.zeros', (['(1)', '(3)', 'imgsz', 'imgsz'], {}), '(1, 3, imgsz, imgsz)\n', (2866, 2886), False, 'import torch\n')]