tracking_countermeasure.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn      \n",
    "import torch.nn.functional as F\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.pyplot import figure\n",
    "import cv2\n",
    "from PIL import Image\n",
    "from pylab import rcParams\n",
    "rcParams['figure.figsize'] = 5, 5\n",
    "import tqdm\n",
    "from numpy import moveaxis\n",
    "import scipy.ndimage\n",
    "from IPython.display import clear_output\n",
    "import time\n",
    "import pyautogui\n",
    "from Xlib import display\n",
    "from mss.linux import MSS as mss\n",
    "import mss.tools\n",
    "\n",
    "animals = [\"chicken\", \"pig\", \"panda\", \"terrain\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def screen_cap(top = 0, left = 0, width = 650, height = 480):\n",
    "    mon = {\"top\": top, \"left\": left, \"width\": width, \"height\": height}\n",
    "    sct = mss.mss()\n",
    "    img = np.array(sct.grab(mon))[...,:3]\n",
    "    \n",
    "    sct.close()\n",
    "    return img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def shoot_arrow():\n",
    "    pyautogui.click(button=\"right\")\n",
    "    \n",
    "def reload_arrow(time_taken):\n",
    "    pyautogui.mouseDown(button = \"right\")\n",
    "    time.sleep(time_taken)\n",
    "    pyautogui.mouseUp(button = \"right\")\n",
    "\n",
    "    \n",
    "def strike_with_sword():\n",
    "    pyautogui.click(button=\"left\")\n",
    "\n",
    "\n",
    "def normalise_targets(target, actual_image_size, small_image_size):\n",
    "    \n",
    "    new_x = int((target[0]/small_image_size[0])*actual_image_size[0])\n",
    "    new_y = int((target[1]/small_image_size[1])*actual_image_size[1])    \n",
    "    return [new_x, new_y]\n",
    "    \n",
    "\n",
    "def move_mouse_co_ord_x(x_val , time = 0.001):\n",
    "    y_val = find_mouse(y = True)\n",
    "    pyautogui.moveTo(x_val,y_val, duration= time )\n",
    "    \n",
    "def move_mouse_co_ord(x_val, y_val , time = 0.001):\n",
    "    pyautogui.moveTo(x_val,y_val, duration= time )\n",
    "    \n",
    "    \n",
    "def find_mouse(x = False , y = False):\n",
    "    \n",
    "    data = display.Display().screen().root.query_pointer()._data\n",
    "    res = [data[\"root_x\"], data[\"root_y\"]]\n",
    "    if x == True and y == True:\n",
    "        return res\n",
    "    elif x == True and y == False:\n",
    "        return res[0]\n",
    "    else:\n",
    "        return res[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Net(\n",
       "  (conv1): Conv2d(3, 7, kernel_size=(5, 5), stride=(1, 1))\n",
       "  (pool): MaxPool2d(kernel_size=3, stride=3, padding=0, dilation=1, ceil_mode=False)\n",
       "  (dropout): Dropout(p=0.5, inplace=False)\n",
       "  (conv2): Conv2d(7, 10, kernel_size=(5, 5), stride=(1, 1))\n",
       "  (fc1): Linear(in_features=10, out_features=10, bias=True)\n",
       "  (fc2): Linear(in_features=10, out_features=8, bias=True)\n",
       "  (fc3): Linear(in_features=8, out_features=4, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(3, 7, 5)\n",
    "        self.pool = nn.MaxPool2d(3, 3)\n",
    "        \n",
    "        self.dropout = nn.Dropout(p=0.5)\n",
    "        self.conv2 = nn.Conv2d(7, 10, 5)\n",
    "        \n",
    "        self.fc1 = nn.Linear(10, 10)\n",
    "        self.fc2 = nn.Linear(10, 8)\n",
    "        self.fc3 = nn.Linear(8, 4)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.pool(F.relu(self.conv1(x)))\n",
    "        x = self.pool(F.relu(self.conv2(x)))\n",
    "        x = x.view(x.size(0), -1)\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = F.relu(self.fc2(x))\n",
    "        x = self.fc3(x)\n",
    "        return x\n",
    "    \n",
    "hunter = Net()\n",
    "hunter.load_state_dict(torch.load(\"hunter.pt\"))\n",
    "hunter.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fgsm(image, epsilon, data_grad):\n",
    "    # Collect the element-wise sign of the data gradient\n",
    "    sign_data_grad = data_grad.sign()\n",
    "        # Create the perturbed image by adjusting each pixel of the input image\n",
    "    perturbed_image = image + epsilon*sign_data_grad\n",
    "        \n",
    "    perturbed_image = torch.clamp(perturbed_image, 0, 255)\n",
    "    \n",
    "    return perturbed_image\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def negative_epsilon_fgsm_on_tensor(input_tensor, model, epsilon, target_label):\n",
    "    \n",
    "    loss = nn.CrossEntropyLoss()\n",
    "    \n",
    "    input_tensor.requires_grad = True  # gradients \n",
    "    label = torch.tensor([target_label], dtype=torch.long)  \n",
    "    pred = model(input_tensor)\n",
    "    loss_val = loss(pred,label)\n",
    "    loss_val.backward() #backprop\n",
    "    data_grad = input_tensor.grad.data\n",
    "    \n",
    "    bamboozle_tensor = fgsm(input_tensor, epsilon, data_grad)\n",
    "   \n",
    "    return bamboozle_tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def return_target_numpy(input_np_arr, model, kernel_size, stride, thresh_prob, attack = False, attack_label = 0):\n",
    "    image = input_np_arr\n",
    "    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "    width, height = image.shape[1], image.shape[0]\n",
    "    circles = []\n",
    "    \n",
    "    final_index = -1\n",
    "    animal_indexes = []\n",
    "    move = False\n",
    "    image_ret = image.copy()\n",
    "    for h in range(0, height):\n",
    "        for w in range(0, width):\n",
    "            h_start = h*stride\n",
    "            w_start = w*stride\n",
    "            h_end = min(height, h_start + kernel_size)\n",
    "            w_end = min(width, w_start + kernel_size)\n",
    "            if (w_end) >= width or (h_end) >= height:\n",
    "                continue\n",
    "            np_slice = image[h_start:h_end, w_start:w_end, :]\n",
    "            if 0 in np_slice.shape:\n",
    "                continue\n",
    "            np_slice = cv2.resize(np_slice, (25,25))   \n",
    "            image_moveaxis =  np.moveaxis(np_slice, 2, 0)    \n",
    "            sweet_tensor =  torch.from_numpy(image_moveaxis).unsqueeze(0).float()\n",
    "            \n",
    "            \n",
    "                \n",
    "            output_dist = model(sweet_tensor).flatten()\n",
    "            output_dist_soft = torch.softmax(output_dist, dim = 0)            \n",
    "            index = torch.argmax(output_dist_soft).item()\n",
    "            prob = output_dist_soft[index].item()\n",
    "            \n",
    "            if index <= 2: ## dont detect terrain\n",
    "                \n",
    "                if index != attack_label:\n",
    "                    if attack == True:\n",
    "                        \n",
    "                        # very bad code here, change later \n",
    "                        sweet_tensor = negative_epsilon_fgsm_on_tensor(sweet_tensor, model, epsilon = -25, target_label = attack_label)\n",
    "\n",
    "                        output_dist = model(sweet_tensor).flatten()\n",
    "                        output_dist_soft = torch.softmax(output_dist, dim = 0)            \n",
    "                        index = torch.argmax(output_dist_soft).item()\n",
    "                        prob = output_dist_soft[index].item()\n",
    "                \n",
    "                \n",
    "                \n",
    "                if  prob > thresh_prob:\n",
    "                    animal_indexes.append(index)\n",
    "                    circle_center = [int((w_end+w_start)/2), int((h_end+h_start)/2)]\n",
    "                    circles.append(circle_center)\n",
    "                else:\n",
    "                    pass\n",
    "    try:        \n",
    "        final_index =  np.bincount(animal_indexes).argmax()\n",
    "\n",
    "    except:\n",
    "        pass\n",
    "    target = np.array([])\n",
    "    if len(circles) != 0:\n",
    "        move = True\n",
    "        target = np.mean(np.array(circles), axis = 0).astype(np.uint8)\n",
    "    \n",
    "    return target, move, final_index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mouse_movement_and_strike(target, top_offset):\n",
    "        try:\n",
    "            # move mouse to a certain pos\n",
    "            move_mouse_co_ord(x_val = target[0] , y_val = target[1] + top_offset, time = 0.0001)\n",
    "            strike_with_sword()\n",
    "        except:\n",
    "            #panic mode basically\n",
    "            \n",
    "            cv2.destroyAllWindows() \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "font = cv2.FONT_HERSHEY_DUPLEX  \n",
    "# fontScale \n",
    "fontScale = 1\n",
    "   \n",
    "# Blue color in BGR \n",
    "color = (0, 255, 255) \n",
    "  \n",
    "# Line thickness of 2 px \n",
    "thickness = 2\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "6  FPS\n",
      "found chicken at =  [74 27]\n"
     ]
    }
   ],
   "source": [
    "#  some time to breathe and pray that it doesn't embarass me by tracking sunflowers instead of chickens\n",
    "time.sleep(3)  \n",
    "\n",
    "attack = False\n",
    "\n",
    "while True:\n",
    "    \n",
    "    top_offset = 150\n",
    "    \n",
    "    center = (460, 180 )  # crosshairs on the screen\n",
    "    start = time.time()\n",
    "    frame = screen_cap(top= top_offset, height = 550 - top_offset, width = 800)\n",
    "    \n",
    "    #resize into a smaller image size \n",
    "    smaller_frame = cv2.resize(frame, (int(frame.shape[1]/6), int(frame.shape[0]/6)))\n",
    "    \n",
    "    \n",
    "    target, move, index = return_target_numpy(smaller_frame ,hunter, 20, 12, 0.99, attack = attack, attack_label = 1)\n",
    "    \n",
    "    #move mouse only when you see a chicken\n",
    "    if move == True:\n",
    "        print(\"found chicken at = \", target)\n",
    "        target = normalise_targets(target, frame.shape[:-1], smaller_frame.shape[:-1])\n",
    "        \n",
    "        # drawing circles and lines because why not, slower but cooler looking pipeline\n",
    "        frame = cv2.circle(frame, tuple(target), radius = 50, color = (0, 100, 255) , thickness = 2)\n",
    "        frame = cv2.line(frame, center, tuple(target), color = (0,0,255), thickness = 5)\n",
    "\n",
    "        \n",
    "        frame = cv2.putText(frame, animals[index], tuple(target), font,fontScale, color, thickness, cv2.LINE_AA) \n",
    "                            \n",
    "                            # weird flex by showing the live feed with circles and lines       \n",
    "    cv2.imshow(\"hunter_fpv\",frame)\n",
    "    \n",
    "    #press q to quit\n",
    "    if cv2.waitKey(1) & 0xFF == ord('q'):\n",
    "        cv2.destroyAllWindows() \n",
    "        break\n",
    "      \n",
    "    # metrics\n",
    "    time_taken = time.time() - start\n",
    "    clear_output(wait = True)\n",
    "\n",
    "    print(int(1/time_taken), \" FPS\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cv2.destroyAllWindows() "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}